WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s

This commit is contained in:
yawan
2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
+265
View File
@@ -0,0 +1,265 @@
package web_search
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"regexp"
"strings"
"time"
"unicode"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const (
// defaultBaiduWebSearchURL is the hardcoded Baidu AI Search API URL.
// Not configurable by tenants — prevents SSRF.
defaultBaiduWebSearchURL = "https://qianfan.baidubce.com/v2/ai_search/web_search"
defaultBaiduTimeout = 15 * time.Second
defaultBaiduResults = 5
maxBaiduResults = 50
maxBaiduQueryUnits = 72
)
// BaiduProvider implements web search using Baidu AI Search API.
type BaiduProvider struct {
client *http.Client
baseURL string
apiKey string
}
// NewBaiduProvider creates a new Baidu web search provider from parameters.
func NewBaiduProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
if params.APIKey == "" {
return nil, fmt.Errorf("API key is required for Baidu provider")
}
client, err := NewSearchHTTPClient(defaultBaiduTimeout, params.ProxyURL)
if err != nil {
return nil, fmt.Errorf("create Baidu HTTP client: %w", err)
}
return &BaiduProvider{
client: client,
baseURL: defaultBaiduWebSearchURL, // Hardcoded — not tenant-configurable
apiKey: params.APIKey,
}, nil
}
// Name returns the provider name.
func (p *BaiduProvider) Name() string {
return "baidu"
}
// Search performs a web search using Baidu AI Search API.
func (p *BaiduProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
preparedQuery := normalizeBaiduQuery(query)
if preparedQuery == "" {
return nil, fmt.Errorf("query is empty")
}
if preparedQuery != strings.TrimSpace(query) {
logger.Infof(ctx, "[WebSearch][Baidu] normalized query to satisfy API constraints")
}
if maxResults <= 0 {
maxResults = defaultBaiduResults
}
if maxResults > maxBaiduResults {
maxResults = maxBaiduResults
}
logger.Infof(ctx, "[WebSearch][Baidu] query=%q maxResults=%d url=%s", preparedQuery, maxResults, p.baseURL)
req, err := p.buildRequest(ctx, preparedQuery, maxResults)
if err != nil {
return nil, err
}
results, err := p.doSearch(ctx, req, includeDate)
if err != nil {
logger.Warnf(ctx, "[WebSearch][Baidu] failed: %v", err)
return nil, err
}
logger.Infof(ctx, "[WebSearch][Baidu] returned %d results", len(results))
return results, nil
}
func (p *BaiduProvider) buildRequest(ctx context.Context, query string, maxResults int) (*http.Request, error) {
requestBody := baiduSearchRequest{
Messages: []baiduMessage{
{Role: "user", Content: query},
},
SearchSource: "baidu_search_v2",
ResourceTypeFilter: []baiduResourceTypeFilter{
{Type: "web", TopK: maxResults},
},
}
bodyBytes, err := json.Marshal(requestBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request body: %w", err)
}
req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL, bytes.NewReader(bodyBytes))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+p.apiKey)
req.Header.Set("Content-Type", "application/json")
return req, nil
}
func (p *BaiduProvider) doSearch(ctx context.Context, req *http.Request, includeDate bool) ([]*types.WebSearchResult, error) {
resp, err := p.client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
body, err := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) // 2MB limit
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
logger.Warnf(ctx, "[WebSearch][Baidu] API returned status %d: %s", resp.StatusCode, string(body))
return nil, fmt.Errorf("baidu API returned status %d: %s", resp.StatusCode, string(body))
}
var respData baiduSearchResponse
if err := json.Unmarshal(body, &respData); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
// Check for API-level error (returned with 200 status)
if respData.Code != 0 {
return nil, fmt.Errorf("baidu API error (code %d): %s", respData.Code, respData.Message)
}
results := make([]*types.WebSearchResult, 0, len(respData.References))
for _, ref := range respData.References {
result := &types.WebSearchResult{
Title: ref.Title,
URL: ref.URL,
Content: ref.Content,
Source: "baidu",
}
if includeDate && ref.Date != "" {
if t, err := parseBaiduDate(ref.Date); err == nil {
result.PublishedAt = &t
}
}
results = append(results, result)
}
return results, nil
}
var baiduDateRe = regexp.MustCompile(`^(\d{4})-(\d{1,2})-(\d{1,2})(?:\s+(\d{1,2}):(\d{2})(?::(\d{2}))?)?`)
// parseBaiduDate extracts date components via regex, handling variable formats
// like "2025-4-24", "2025-04-27 18:02:00", "2025-05-20 11:58" uniformly.
func parseBaiduDate(dateStr string) (time.Time, error) {
m := baiduDateRe.FindStringSubmatch(dateStr)
if m == nil {
return time.Time{}, fmt.Errorf("unable to parse date: %s", dateStr)
}
// Pad to "YYYY-MM-DD HH:MM:SS" and parse once
normalized := fmt.Sprintf("%s-%02s-%02s %02s:%02s:%02s",
m[1], m[2], m[3],
defaultStr(m[4], "00"), defaultStr(m[5], "00"), defaultStr(m[6], "00"))
return time.Parse("2006-01-02 15:04:05", normalized)
}
func defaultStr(s, fallback string) string {
if s == "" {
return fallback
}
return s
}
// Baidu documents the query length limit as 72 chars, counting CJK/full-width
// runes as 2. Use a conservative width model so mixed-language input stays
// under the API limit.
func normalizeBaiduQuery(query string) string {
query = strings.TrimSpace(query)
if query == "" {
return ""
}
if baiduQueryUnits(query) <= maxBaiduQueryUnits {
return query
}
var b strings.Builder
b.Grow(len(query))
used := 0
for _, r := range query {
width := baiduQueryUnitWidth(r)
if used+width > maxBaiduQueryUnits {
break
}
b.WriteRune(r)
used += width
}
return b.String()
}
func baiduQueryUnits(query string) int {
units := 0
for _, r := range query {
units += baiduQueryUnitWidth(r)
}
return units
}
func baiduQueryUnitWidth(r rune) int {
if r <= unicode.MaxASCII {
return 1
}
return 2
}
// --- Request/Response types ---
type baiduSearchRequest struct {
Messages []baiduMessage `json:"messages"`
SearchSource string `json:"search_source"`
ResourceTypeFilter []baiduResourceTypeFilter `json:"resource_type_filter"`
SearchRecencyFilter string `json:"search_recency_filter,omitempty"`
}
type baiduMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
type baiduResourceTypeFilter struct {
Type string `json:"type"`
TopK int `json:"top_k"`
}
type baiduSearchResponse struct {
References []baiduReference `json:"references"`
RequestID string `json:"request_id"`
Code int `json:"code"`
Message string `json:"message"`
}
type baiduReference struct {
ID int `json:"id"`
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
Date string `json:"date"`
Type string `json:"type"`
}
+199
View File
@@ -0,0 +1,199 @@
package web_search
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const (
// defaultBingSearchURL is the hardcoded Bing search API URL.
// Not configurable by tenants — prevents SSRF.
defaultBingSearchURL = "https://api.bing.microsoft.com/v7.0/search"
)
var (
defaultUserAgentHeader = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36"
defaultBingTimeout = 10 * time.Second
)
type bingSafeSearch string
const (
bingSafeSearchOff bingSafeSearch = "Off"
bingSafeSearchModerate bingSafeSearch = "Moderate"
bingSafeSearchStrict bingSafeSearch = "Strict"
)
type bingFreshness string
const (
bingFreshnessDay = "Day"
bingFreshnessWeek = "Week"
bingFreshnessMonth = "Month"
)
// BingProvider implements web search using Bing Search API
type BingProvider struct {
client *http.Client
baseURL string
apiKey string
}
// NewBingProvider creates a new Bing provider from parameters (no environment variables).
func NewBingProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
if params.APIKey == "" {
return nil, fmt.Errorf("API key is required for Bing provider")
}
client, err := NewSearchHTTPClient(defaultBingTimeout, params.ProxyURL)
if err != nil {
return nil, err
}
return &BingProvider{
client: client,
baseURL: defaultBingSearchURL, // Hardcoded — not tenant-configurable
apiKey: params.APIKey,
}, nil
}
// Name returns the provider name
func (p *BingProvider) Name() string {
return "bing"
}
// Search performs a web search using Bing Search API
func (p *BingProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
if len(query) == 0 {
return nil, fmt.Errorf("query is empty")
}
logger.Infof(ctx, "[WebSearch][Bing] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
req, err := p.buildParams(ctx, query, maxResults, includeDate)
if err != nil {
return nil, err
}
results, err := p.doSearch(ctx, req)
if err != nil {
logger.Warnf(ctx, "[WebSearch][Bing] failed: %v", err)
return nil, err
}
logger.Infof(ctx, "[WebSearch][Bing] returned %d results", len(results))
return results, nil
}
func (p *BingProvider) doSearch(ctx context.Context, req *http.Request) ([]*types.WebSearchResult, error) {
resp, err := p.client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
logger.Warnf(ctx, "[WebSearch][Bing] API returned status %d: %s", resp.StatusCode, string(body))
return nil, fmt.Errorf("bing API returned status %d: %s", resp.StatusCode, string(body))
}
var respData bingSearchResponse
if err := json.Unmarshal(body, &respData); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
results := make([]*types.WebSearchResult, 0, len(respData.WebPages.Value))
for _, item := range respData.WebPages.Value {
results = append(results, &types.WebSearchResult{
Title: item.Name,
URL: item.URL,
Snippet: item.Snippet,
Source: "bing",
PublishedAt: &item.DateLastCrawled,
})
}
return results, nil
}
// bingSearchResponse defines the response structure for Bing search API.
type bingSearchResponse struct {
Type string `json:"_type"`
QueryContext struct {
OriginalQuery string `json:"originalQuery"`
} `json:"queryContext"`
WebPages struct {
WebSearchURL string `json:"webSearchUrl"`
TotalEstimatedMatches int `json:"totalEstimatedMatches"`
Value []struct {
ID string `json:"id"`
Name string `json:"name"`
URL string `json:"url"`
IsFamilyFriendly bool `json:"isFamilyFriendly"`
DisplayURL string `json:"displayUrl"`
Snippet string `json:"snippet"`
DateLastCrawled time.Time `json:"dateLastCrawled"`
SearchTags []struct {
Name string `json:"name"`
Content string `json:"content"`
} `json:"searchTags,omitempty"`
About []struct {
Name string `json:"name"`
} `json:"about,omitempty"`
} `json:"value"`
} `json:"webPages"`
RelatedSearches struct {
ID string `json:"id"`
Value []struct {
Text string `json:"text"`
DisplayText string `json:"displayText"`
WebSearchURL string `json:"webSearchUrl"`
} `json:"value"`
} `json:"relatedSearches"`
RankingResponse struct {
Mainline struct {
Items []struct {
AnswerType string `json:"answerType"`
ResultIndex int `json:"resultIndex"`
Value struct {
ID string `json:"id"`
} `json:"value"`
} `json:"items"`
} `json:"mainline"`
Sidebar struct {
Items []struct {
AnswerType string `json:"answerType"`
Value struct {
ID string `json:"id"`
} `json:"value"`
} `json:"items"`
} `json:"sidebar"`
} `json:"rankingResponse"`
}
func (p *BingProvider) buildParams(ctx context.Context, query string, maxResults int, includeDate bool) (*http.Request, error) {
params := url.Values{}
params.Set("q", query)
params.Set("count", strconv.Itoa(maxResults))
queryURL := fmt.Sprintf("%s?%s", p.baseURL, params.Encode())
req, err := http.NewRequestWithContext(ctx, "GET", queryURL, nil)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("User-Agent", defaultUserAgentHeader)
req.Header.Set("Ocp-Apim-Subscription-Key", p.apiKey)
return req, nil
}
@@ -0,0 +1,254 @@
package web_search
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"time"
"github.com/PuerkitoBio/goquery"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// DuckDuckGoProvider implements web search using DuckDuckGo (HTML first, API fallback)
type DuckDuckGoProvider struct {
client *http.Client
}
// NewDuckDuckGoProvider creates a new DuckDuckGo provider.
// DuckDuckGo is free and requires no API key; optional ProxyURL tunnels requests via an HTTP/HTTPS proxy.
func NewDuckDuckGoProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
client, err := NewSearchHTTPClient(30*time.Second, params.ProxyURL)
if err != nil {
return nil, err
}
return &DuckDuckGoProvider{client: client}, nil
}
// Name returns the provider name
func (p *DuckDuckGoProvider) Name() string {
return "duckduckgo"
}
// Search performs a web search using DuckDuckGo HTML endpoint with API fallback
func (p *DuckDuckGoProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
if maxResults <= 0 {
maxResults = 5
}
// Try HTML scraping first (more reliable for general results)
htmlResults, err := p.searchHTML(ctx, query, maxResults)
if err == nil && len(htmlResults) > 0 {
return htmlResults, nil
}
// Fallback to Instant Answer API
apiResults, apiErr := p.searchAPI(ctx, query, maxResults)
if apiErr == nil && len(apiResults) > 0 {
return apiResults, nil
}
if err != nil {
return nil, fmt.Errorf("duckduckgo HTML search failed: %w", err)
}
return nil, fmt.Errorf("duckduckgo API search failed: %w", apiErr)
}
// searchHTML performs a web search using DuckDuckGo HTML endpoint
func (p *DuckDuckGoProvider) searchHTML(
ctx context.Context,
query string,
maxResults int,
) ([]*types.WebSearchResult, error) {
baseURL := "https://html.duckduckgo.com/html/"
params := url.Values{}
params.Set("q", query)
params.Set("kl", "cn-zh")
reqURL := baseURL + "?" + params.Encode()
req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set(
"User-Agent",
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
)
curlCommand := fmt.Sprintf(
"curl -X GET '%s' -H 'User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'",
req.URL.String(),
)
logger.Infof(ctx, "Curl of request: %s", secutils.SanitizeForLog(curlCommand))
resp, err := p.client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to perform request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
return nil, fmt.Errorf("duckduckgo HTML returned status %d", resp.StatusCode)
}
doc, err := goquery.NewDocumentFromReader(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to parse HTML: %w", err)
}
results := make([]*types.WebSearchResult, 0, maxResults)
doc.Find(".web-result").Each(func(i int, s *goquery.Selection) {
if len(results) >= maxResults {
return
}
titleNode := s.Find(".result__a")
title := strings.TrimSpace(titleNode.Text())
var link string
if href, exists := titleNode.Attr("href"); exists {
link = cleanDDGURL(href)
}
snippet := strings.TrimSpace(s.Find(".result__snippet").Text())
if title != "" && link != "" {
results = append(results, &types.WebSearchResult{
Title: title,
URL: link,
Snippet: snippet,
Source: "duckduckgo",
})
}
})
logger.Infof(ctx, "DuckDuckGo HTML search returned %d results for query: %s", len(results), query)
return results, nil
}
// searchAPI performs a web search using DuckDuckGo API endpoint
func (p *DuckDuckGoProvider) searchAPI(
ctx context.Context,
query string,
maxResults int,
) ([]*types.WebSearchResult, error) {
baseURL := "https://api.duckduckgo.com/"
params := url.Values{}
params.Set("q", query)
params.Set("format", "json")
params.Set("no_html", "1")
params.Set("skip_disambig", "1")
reqURL := baseURL + "?" + params.Encode()
req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("User-Agent", "WeKnora/1.0")
resp, err := p.client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to perform request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("duckduckgo API returned status %d: %s", resp.StatusCode, string(body))
}
var apiResponse struct {
AbstractText string `json:"AbstractText"`
AbstractURL string `json:"AbstractURL"`
Heading string `json:"Heading"`
RelatedTopics []struct {
FirstURL string `json:"FirstURL"`
Text string `json:"Text"`
} `json:"RelatedTopics"`
Results []struct {
FirstURL string `json:"FirstURL"`
Text string `json:"Text"`
} `json:"Results"`
}
if err := json.NewDecoder(resp.Body).Decode(&apiResponse); err != nil {
return nil, fmt.Errorf("failed to decode API response: %w", err)
}
results := make([]*types.WebSearchResult, 0, maxResults)
if apiResponse.AbstractText != "" && apiResponse.AbstractURL != "" {
results = append(results, &types.WebSearchResult{
Title: apiResponse.Heading,
URL: apiResponse.AbstractURL,
Snippet: apiResponse.AbstractText,
Source: "duckduckgo",
})
}
for _, topic := range apiResponse.RelatedTopics {
if len(results) >= maxResults {
break
}
if topic.Text != "" && topic.FirstURL != "" {
results = append(results, &types.WebSearchResult{
Title: extractTitle(topic.Text),
URL: topic.FirstURL,
Snippet: topic.Text,
Source: "duckduckgo",
})
}
}
for _, r := range apiResponse.Results {
if len(results) >= maxResults {
break
}
if r.Text != "" && r.FirstURL != "" {
results = append(results, &types.WebSearchResult{
Title: extractTitle(r.Text),
URL: r.FirstURL,
Snippet: r.Text,
Source: "duckduckgo",
})
}
}
logger.Infof(ctx, "DuckDuckGo API search returned %d results for query: %s", len(results), query)
return results, nil
}
// cleanDDGURL cleans the URL from DuckDuckGo HTML endpoint
func cleanDDGURL(urlStr string) string {
if strings.HasPrefix(urlStr, "//duckduckgo.com/l/?uddg=") {
trimmed := strings.TrimPrefix(urlStr, "//duckduckgo.com/l/?uddg=")
if idx := strings.Index(trimmed, "&rut="); idx != -1 {
decodedStr, err := url.PathUnescape(trimmed[:idx])
if err == nil {
return decodedStr
}
return ""
}
}
if strings.HasPrefix(urlStr, "https://duckduckgo.com/l/?uddg=") {
if parsedURL, err := url.Parse(urlStr); err == nil {
if uddg := parsedURL.Query().Get("uddg"); uddg != "" {
return uddg
}
}
}
return urlStr
}
// extractTitle extracts the title from the text
func extractTitle(text string) string {
lines := strings.Split(text, "\n")
if len(lines) > 0 {
title := strings.TrimSpace(lines[0])
if len(title) > 100 {
title = title[:100] + "..."
}
return title
}
return strings.TrimSpace(text)
}
+181
View File
@@ -0,0 +1,181 @@
package web_search
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const (
defaultExaSearchURL = "https://api.exa.ai/search"
defaultExaTimeout = 15 * time.Second
defaultExaResults = 5
maxExaResults = 100
maxExaResponseBytes = 2 << 20
maxExaContentRunes = 12000
)
// ExaProvider implements web search using Exa's official Search API.
type ExaProvider struct {
client *http.Client
baseURL string
apiKey string
includeText bool
}
// NewExaProvider creates an Exa provider from tenant-specific parameters.
func NewExaProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
apiKey := strings.TrimSpace(params.APIKey)
if apiKey == "" {
return nil, fmt.Errorf("API key is required for Exa provider")
}
client, err := NewSearchHTTPClient(defaultExaTimeout, params.ProxyURL)
if err != nil {
return nil, err
}
return &ExaProvider{
client: client,
baseURL: defaultExaSearchURL,
apiKey: apiKey,
includeText: parseExaBool(params.ExtraConfig, "include_text"),
}, nil
}
// Name returns the provider type identifier.
func (p *ExaProvider) Name() string { return "exa" }
// Search performs a web search through Exa's official Search API.
func (p *ExaProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
query = strings.TrimSpace(query)
if query == "" {
return nil, fmt.Errorf("query is empty")
}
if maxResults <= 0 {
maxResults = defaultExaResults
}
if maxResults > maxExaResults {
maxResults = maxExaResults
}
bodyBytes, err := json.Marshal(exaSearchRequest{
Query: query,
NumResults: maxResults,
Contents: exaContents{
Highlights: true,
Text: p.includeText,
},
})
if err != nil {
return nil, fmt.Errorf("failed to marshal Exa request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL, bytes.NewReader(bodyBytes))
if err != nil {
return nil, fmt.Errorf("failed to create Exa request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("x-api-key", p.apiKey)
logger.Infof(ctx, "[WebSearch][Exa] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
resp, err := p.client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute Exa request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(io.LimitReader(resp.Body, maxExaResponseBytes))
if err != nil {
return nil, fmt.Errorf("failed to read Exa response: %w", err)
}
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
logger.Warnf(ctx, "[WebSearch][Exa] API returned status %d: %s", resp.StatusCode, string(body))
return nil, fmt.Errorf("exa API returned status %d: %s", resp.StatusCode, string(body))
}
var data exaSearchResponse
if err := json.Unmarshal(body, &data); err != nil {
return nil, fmt.Errorf("failed to unmarshal Exa response: %w", err)
}
if data.Error != "" {
return nil, fmt.Errorf("exa API error: %s", data.Error)
}
results := make([]*types.WebSearchResult, 0, len(data.Results))
for _, item := range data.Results {
if len(results) >= maxResults {
break
}
snippet := strings.TrimSpace(strings.Join(item.Highlights, "\n"))
content := truncateExaText(strings.TrimSpace(item.Text), maxExaContentRunes)
if snippet == "" {
snippet = truncateExaText(content, 500)
}
result := &types.WebSearchResult{Title: item.Title, URL: item.URL, Snippet: snippet, Content: content, Source: "exa"}
if includeDate && item.PublishedDate != "" {
if publishedAt, err := time.Parse(time.RFC3339, item.PublishedDate); err == nil {
result.PublishedAt = &publishedAt
}
}
results = append(results, result)
}
logger.Infof(ctx, "[WebSearch][Exa] returned %d results", len(results))
return results, nil
}
func parseExaBool(config map[string]string, key string) bool {
v, err := strconv.ParseBool(strings.TrimSpace(config[key]))
return err == nil && v
}
func truncateExaText(value string, maxRunes int) string {
if maxRunes <= 0 {
return ""
}
if utf8.RuneCountInString(value) <= maxRunes {
return value
}
byteIndex := 0
for runeIndex := 0; runeIndex < maxRunes && byteIndex < len(value); runeIndex++ {
_, size := utf8.DecodeRuneInString(value[byteIndex:])
byteIndex += size
}
return value[:byteIndex]
}
type exaSearchRequest struct {
Query string `json:"query"`
NumResults int `json:"numResults"`
Contents exaContents `json:"contents"`
}
type exaContents struct {
Highlights bool `json:"highlights"`
Text bool `json:"text,omitempty"`
}
type exaSearchResponse struct {
Results []exaResult `json:"results"`
Error string `json:"error,omitempty"`
}
type exaResult struct {
Title string `json:"title"`
URL string `json:"url"`
PublishedDate string `json:"publishedDate,omitempty"`
Highlights []string `json:"highlights,omitempty"`
Text string `json:"text,omitempty"`
}
@@ -0,0 +1,117 @@
package web_search
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
)
func TestExaProviderSearchMapping(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost || r.URL.Path != "/search" {
t.Fatalf("request = %s %s", r.Method, r.URL.Path)
}
if got := r.Header.Get("x-api-key"); got != "exa-test" {
t.Fatalf("x-api-key = %q", got)
}
var request exaSearchRequest
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Fatal(err)
}
if request.Query != "hello" || request.NumResults != 2 ||
!request.Contents.Highlights || !request.Contents.Text {
t.Fatalf("request body = %+v", request)
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(exaSearchResponse{Results: []exaResult{
{
Title: "One",
URL: "https://example.com/1",
Highlights: []string{"first", "second"},
Text: "body",
PublishedDate: "2026-05-01T00:00:00Z",
},
{Title: "Two", URL: "https://example.com/2"},
{Title: "Three", URL: "https://example.com/3"},
}})
}))
defer srv.Close()
p := &ExaProvider{client: srv.Client(), baseURL: srv.URL + "/search", apiKey: "exa-test", includeText: true}
results, err := p.Search(context.Background(), "hello", 2, true)
if err != nil {
t.Fatal(err)
}
if len(results) != 2 || results[0].Snippet != "first\nsecond" ||
results[0].Content != "body" || results[0].Source != "exa" {
t.Fatalf("unexpected results: %+v", results)
}
if results[0].PublishedAt == nil || !results[0].PublishedAt.Equal(time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC)) {
t.Fatalf("unexpected date: %v", results[0].PublishedAt)
}
if results[1].Snippet != "" {
t.Fatalf("unexpected empty-result fallback: %+v", results[1])
}
}
func TestExaProviderSupportsConnectionTestProbe(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost || r.URL.Path != "/search" {
t.Fatalf("request = %s %s", r.Method, r.URL.Path)
}
if got := r.Header.Get("x-api-key"); got != "exa-test" {
t.Fatalf("x-api-key = %q", got)
}
var request exaSearchRequest
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Fatal(err)
}
if request.Query != "test" || request.NumResults != 1 {
t.Fatalf("connection test request = %+v", request)
}
if !request.Contents.Highlights || request.Contents.Text {
t.Fatalf("connection test should request highlights without page text: %+v", request.Contents)
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(exaSearchResponse{Results: []exaResult{
{Title: "Exa", URL: "https://exa.ai/"},
}})
}))
defer srv.Close()
p := &ExaProvider{
client: srv.Client(),
baseURL: srv.URL + "/search",
apiKey: "exa-test",
}
results, err := p.Search(context.Background(), "test", 1, false)
if err != nil {
t.Fatal(err)
}
if len(results) != 1 {
t.Fatalf("len(results) = %d, want 1", len(results))
}
}
func TestExaProviderValidationAndStatus(t *testing.T) {
if _, err := NewExaProvider(types.WebSearchProviderParameters{}); err == nil {
t.Fatal("expected missing API key error")
}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusTooManyRequests)
_, _ = w.Write([]byte(`{"error":"rate limited"}`))
}))
defer srv.Close()
p := &ExaProvider{client: srv.Client(), baseURL: srv.URL, apiKey: "key"}
if _, err := p.Search(context.Background(), "q", 1, false); err == nil {
t.Fatal("expected status error")
}
if _, err := p.Search(context.Background(), " ", 1, false); err == nil {
t.Fatal("expected empty query error")
}
}
@@ -0,0 +1,94 @@
package web_search
import (
"context"
"fmt"
"time"
"google.golang.org/api/customsearch/v1"
"google.golang.org/api/option"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
// GoogleProvider implements web search using Google Custom Search Engine API
type GoogleProvider struct {
srv *customsearch.Service
apiKey string
engineID string
}
// NewGoogleProvider creates a new Google provider from parameters (no environment variables).
// The API endpoint is the official Google Custom Search endpoint — not tenant-configurable.
func NewGoogleProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
if params.APIKey == "" {
return nil, fmt.Errorf("API key is required for Google provider")
}
if params.EngineID == "" {
return nil, fmt.Errorf("engine ID is required for Google provider")
}
httpClient, err := NewSearchHTTPClient(30*time.Second, params.ProxyURL)
if err != nil {
return nil, err
}
clientOpts := []option.ClientOption{
option.WithAPIKey(params.APIKey),
option.WithHTTPClient(httpClient),
}
srv, err := customsearch.NewService(context.Background(), clientOpts...)
if err != nil {
return nil, err
}
return &GoogleProvider{
srv: srv,
apiKey: params.APIKey,
engineID: params.EngineID,
}, nil
}
// Name returns the provider name
func (p *GoogleProvider) Name() string {
return "google"
}
// Search performs a web search using Google Custom Search Engine API
func (p *GoogleProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
if len(query) == 0 {
return nil, fmt.Errorf("query is empty")
}
logger.Infof(ctx, "[WebSearch][Google] query=%q maxResults=%d engineID=%s", query, maxResults, p.engineID)
cseCall := p.srv.Cse.List().Context(ctx).Cx(p.engineID).Q(query)
if maxResults > 0 {
cseCall = cseCall.Num(int64(maxResults))
} else {
cseCall = cseCall.Num(5)
}
cseCall = cseCall.Hl("ch-zh")
resp, err := cseCall.Do()
if err != nil {
logger.Warnf(ctx, "[WebSearch][Google] failed: %v", err)
return nil, err
}
results := make([]*types.WebSearchResult, 0)
for _, item := range resp.Items {
result := &types.WebSearchResult{
Title: item.Title,
URL: item.Link,
Snippet: item.Snippet,
Source: "google",
}
results = append(results, result)
}
logger.Infof(ctx, "[WebSearch][Google] returned %d results", len(results))
return results, nil
}
@@ -0,0 +1,168 @@
package web_search
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const (
// defaultKeenableBaseURL is the hardcoded Keenable API base URL.
// Not configurable by tenants — prevents SSRF.
defaultKeenableBaseURL = "https://api.keenable.ai"
// keenableTitle is the attribution tag Keenable segments integration traffic by.
keenableTitle = "WeKnora"
// defaultKeenableResults matches the fallback other providers use when the
// caller does not specify a positive maxResults.
defaultKeenableResults = 5
)
var (
defaultKeenableTimeout = 15 * time.Second
)
// KeenableProvider implements web search using the Keenable Search API.
// Keyless by default: with no API key it calls the public endpoint
// (rate-limited); a key switches to the authenticated endpoint and lifts the cap.
type KeenableProvider struct {
client *http.Client
baseURL string
apiKey string
}
// NewKeenableProvider creates a new Keenable provider from parameters.
// The API key is optional; Keenable works keyless against the public endpoint.
func NewKeenableProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
client, err := NewSearchHTTPClient(defaultKeenableTimeout, params.ProxyURL)
if err != nil {
return nil, err
}
return &KeenableProvider{
client: client,
baseURL: defaultKeenableBaseURL,
apiKey: params.APIKey,
}, nil
}
// Name returns the provider name
func (p *KeenableProvider) Name() string {
return "keenable"
}
// Search performs a web search using the Keenable Search API.
func (p *KeenableProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
if len(query) == 0 {
return nil, fmt.Errorf("query is empty")
}
if maxResults <= 0 {
maxResults = defaultKeenableResults
}
// Keyless by default; a configured key switches to the authenticated path.
path := "/v1/search/public"
if p.apiKey != "" {
path = "/v1/search"
}
endpoint := p.baseURL + path
logger.Infof(ctx, "[WebSearch][Keenable] query=%q maxResults=%d url=%s", query, maxResults, endpoint)
bodyBytes, err := json.Marshal(keenableSearchRequest{Query: query, Mode: "pro"})
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(bodyBytes))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
req.Header.Set("X-Keenable-Title", keenableTitle)
if p.apiKey != "" {
req.Header.Set("X-API-Key", p.apiKey)
}
resp, err := p.client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
respBody, _ := io.ReadAll(resp.Body)
logger.Warnf(ctx, "[WebSearch][Keenable] API returned status %d: %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("keenable API returned status %d: %s", resp.StatusCode, string(respBody))
}
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
var respData keenableSearchResponse
if err := json.Unmarshal(respBody, &respData); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
results := make([]*types.WebSearchResult, 0, len(respData.Results))
for _, item := range respData.Results {
if maxResults > 0 && len(results) >= maxResults {
break
}
// Keenable returns both a short "description" and a longer "snippet"
// excerpt. Prefer the short description as the summary snippet and keep
// the longer excerpt as Content (used by RAG compression). Fall back to
// whichever is present so we never drop the only available text.
snippet := item.Description
if snippet == "" {
snippet = item.Snippet
}
result := &types.WebSearchResult{
Title: item.Title,
URL: item.URL,
Snippet: snippet,
Content: item.Snippet,
Source: "keenable",
}
if includeDate && item.PublishedAt != "" {
if t, err := time.Parse(time.RFC3339, item.PublishedAt); err == nil {
result.PublishedAt = &t
}
}
results = append(results, result)
}
logger.Infof(ctx, "[WebSearch][Keenable] returned %d results", len(results))
return results, nil
}
// keenableSearchRequest defines the request body for the Keenable search API.
type keenableSearchRequest struct {
Query string `json:"query"`
Mode string `json:"mode"`
}
// keenableSearchResponse defines the response structure for the Keenable search API.
type keenableSearchResponse struct {
Query string `json:"query"`
Results []struct {
Title string `json:"title"`
URL string `json:"url"`
Description string `json:"description"`
Snippet string `json:"snippet,omitempty"`
PublishedAt string `json:"published_at,omitempty"`
AcquiredAt string `json:"acquired_at,omitempty"`
} `json:"results"`
}
@@ -0,0 +1,142 @@
package web_search
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
)
func TestKeenableProvider_Search_Keyless(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// No key configured => keyless public endpoint, no X-API-Key.
if r.URL.Path != "/v1/search/public" {
t.Errorf("path = %q, want /v1/search/public", r.URL.Path)
}
if got := r.Header.Get("X-API-Key"); got != "" {
t.Errorf("X-API-Key = %q, want empty for keyless", got)
}
if got := r.Header.Get("X-Keenable-Title"); got != "WeKnora" {
t.Errorf("X-Keenable-Title = %q, want WeKnora", got)
}
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["query"] != "hello" || body["mode"] != "pro" {
t.Errorf("body = %v, want query=hello mode=pro", body)
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"query": "hello",
"results": []map[string]any{
{"title": "T1", "url": "https://e/1", "description": "d1", "published_at": "2026-05-01T00:00:00Z"},
{"title": "T2", "url": "https://e/2", "description": "d2"},
{"title": "T3", "url": "https://e/3", "description": "d3"},
},
})
}))
defer srv.Close()
p := &KeenableProvider{client: srv.Client(), baseURL: srv.URL}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
results, err := p.Search(ctx, "hello", 2, true)
if err != nil {
t.Fatalf("Search: %v", err)
}
if len(results) != 2 {
t.Fatalf("expected 2 results (maxResults trim), got %d", len(results))
}
if results[0].Source != "keenable" {
t.Errorf("source = %q, want keenable", results[0].Source)
}
if results[0].Snippet != "d1" || results[0].URL != "https://e/1" {
t.Errorf("unexpected first result: %+v", results[0])
}
if results[0].PublishedAt == nil {
t.Errorf("expected PublishedAt to be set on the first result")
}
}
func TestKeenableProvider_Search_SnippetMapping(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"results": []map[string]any{
// Has both: description is the short summary, snippet the longer excerpt.
{"title": "T1", "url": "https://e/1", "description": "short", "snippet": "long excerpt"},
// Only snippet present: it must fall back into Snippet so we don't drop text.
{"title": "T2", "url": "https://e/2", "snippet": "only long"},
},
})
}))
defer srv.Close()
p := &KeenableProvider{client: srv.Client(), baseURL: srv.URL}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
results, err := p.Search(ctx, "q", 5, false)
if err != nil {
t.Fatalf("Search: %v", err)
}
if len(results) != 2 {
t.Fatalf("expected 2 results, got %d", len(results))
}
if results[0].Snippet != "short" || results[0].Content != "long excerpt" {
t.Errorf("result[0] snippet/content mapping wrong: %+v", results[0])
}
if results[1].Snippet != "only long" {
t.Errorf("result[1] should fall back to snippet when description is empty: %+v", results[1])
}
}
func TestKeenableProvider_Search_Keyed(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// A configured key => authenticated endpoint + X-API-Key.
if r.URL.Path != "/v1/search" {
t.Errorf("path = %q, want /v1/search", r.URL.Path)
}
if got := r.Header.Get("X-API-Key"); got != "keen_test" {
t.Errorf("X-API-Key = %q, want keen_test", got)
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{"results": []map[string]any{}})
}))
defer srv.Close()
p := &KeenableProvider{client: srv.Client(), baseURL: srv.URL, apiKey: "keen_test"}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if _, err := p.Search(ctx, "hi", 5, false); err != nil {
t.Fatalf("Search: %v", err)
}
}
func TestKeenableProvider_Search_ErrorStatus(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusTooManyRequests)
_, _ = w.Write([]byte(`{"message":"rate limited"}`))
}))
defer srv.Close()
p := &KeenableProvider{client: srv.Client(), baseURL: srv.URL}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if _, err := p.Search(ctx, "x", 5, false); err == nil {
t.Fatal("expected an error for a non-200 response")
}
}
func TestKeenableProvider_Name(t *testing.T) {
p, err := NewKeenableProvider(types.WebSearchProviderParameters{})
if err != nil {
t.Fatalf("NewKeenableProvider (keyless): %v", err)
}
if p.Name() != "keenable" {
t.Errorf("Name() = %q, want keenable", p.Name())
}
}
@@ -0,0 +1,207 @@
package web_search
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const (
defaultMetasoSearchURL = "https://metaso.cn/api/v1/search"
defaultMetasoTimeout = 30 * time.Second
defaultMetasoResults = 10
maxMetasoResults = 50
maxMetasoResponseBytes = 4 << 20
defaultMetasoScope = "webpage"
)
var validMetasoScopes = map[string]struct{}{
"webpage": {}, "document": {}, "scholar": {},
"podcast": {}, "video": {}, "image": {},
}
// MetasoProvider implements web search using the official Metaso AI Search API.
type MetasoProvider struct {
client *http.Client
baseURL string
apiKey string
scope string
}
func NewMetasoProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
if err := ValidateMetasoParameters(params); err != nil {
return nil, err
}
client, err := NewSearchHTTPClient(defaultMetasoTimeout, params.ProxyURL)
if err != nil {
return nil, err
}
return &MetasoProvider{
client: client, baseURL: defaultMetasoSearchURL,
apiKey: strings.TrimSpace(params.APIKey), scope: metasoScope(params.ExtraConfig),
}, nil
}
func ValidateMetasoParameters(params types.WebSearchProviderParameters) error {
if strings.TrimSpace(params.APIKey) == "" {
return fmt.Errorf("API key is required for Metaso provider")
}
scope := metasoScope(params.ExtraConfig)
if _, ok := validMetasoScopes[scope]; !ok {
return fmt.Errorf("invalid Metaso search scope: %s", scope)
}
return nil
}
func metasoScope(extraConfig map[string]string) string {
if scope := strings.TrimSpace(extraConfig["scope"]); scope != "" {
return scope
}
return defaultMetasoScope
}
func (p *MetasoProvider) Name() string { return "metaso" }
func (p *MetasoProvider) Search(ctx context.Context, query string, maxResults int, includeDate bool) ([]*types.WebSearchResult, error) {
query = strings.TrimSpace(query)
if query == "" {
return nil, fmt.Errorf("query is empty")
}
if maxResults <= 0 {
maxResults = defaultMetasoResults
}
if maxResults > maxMetasoResults {
maxResults = maxMetasoResults
}
body, err := json.Marshal(metasoSearchRequest{
Query: query, Scope: p.scope, Size: maxResults,
IncludeSummary: true, IncludeRawContent: false, ConciseSnippet: true,
})
if err != nil {
return nil, fmt.Errorf("failed to marshal Metaso request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("failed to create Metaso request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+p.apiKey)
req.Header.Set("Accept", "application/json")
req.Header.Set("Content-Type", "application/json")
logger.Infof(ctx, "[WebSearch][Metaso] query=%q maxResults=%d scope=%s", query, maxResults, p.scope)
resp, err := p.client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute Metaso request: %w", err)
}
defer resp.Body.Close()
respBody, err := readMetasoResponseBody(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
return nil, metasoHTTPError(resp.StatusCode, respBody)
}
var response metasoSearchResponse
if err := json.Unmarshal(respBody, &response); err != nil {
return nil, fmt.Errorf("failed to unmarshal Metaso response: %w", err)
}
results := make([]*types.WebSearchResult, 0, len(response.Webpages))
for _, item := range response.Webpages {
if strings.TrimSpace(item.Title) == "" && strings.TrimSpace(item.Link) == "" {
continue
}
snippet := strings.TrimSpace(item.Summary)
if snippet == "" {
snippet = strings.TrimSpace(item.Snippet)
}
result := &types.WebSearchResult{Title: item.Title, URL: item.Link, Snippet: snippet, Content: item.RawContent, Source: "metaso"}
if includeDate {
if publishedAt, ok := parseMetasoDate(item.Date); ok {
result.PublishedAt = &publishedAt
}
}
results = append(results, result)
if len(results) >= maxResults {
break
}
}
logger.Infof(ctx, "[WebSearch][Metaso] returned %d results", len(results))
return results, nil
}
func readMetasoResponseBody(reader io.Reader) ([]byte, error) {
body, err := io.ReadAll(io.LimitReader(reader, maxMetasoResponseBytes+1))
if err != nil {
return nil, fmt.Errorf("failed to read Metaso response: %w", err)
}
if len(body) > maxMetasoResponseBytes {
return nil, fmt.Errorf("Metaso response exceeds %d bytes", maxMetasoResponseBytes)
}
return body, nil
}
func metasoHTTPError(statusCode int, body []byte) error {
var apiError struct {
Message string `json:"message"`
Error string `json:"error"`
}
if json.Unmarshal(body, &apiError) == nil {
detail := strings.TrimSpace(apiError.Message)
if detail == "" {
detail = strings.TrimSpace(apiError.Error)
}
if detail != "" {
return fmt.Errorf("Metaso API returned status %d: %s", statusCode, detail)
}
}
detail := strings.TrimSpace(string(body))
if len(detail) > 4096 {
detail = detail[:4096]
}
if detail == "" {
return fmt.Errorf("Metaso API returned status %d", statusCode)
}
return fmt.Errorf("Metaso API returned status %d: %s", statusCode, detail)
}
func parseMetasoDate(value string) (time.Time, bool) {
for _, layout := range []string{time.RFC3339Nano, "2006-01-02 15:04:05", "2006-01-02"} {
if parsed, err := time.Parse(layout, strings.TrimSpace(value)); err == nil {
return parsed, true
}
}
return time.Time{}, false
}
type metasoSearchRequest struct {
Query string `json:"q"`
Scope string `json:"scope"`
Size int `json:"size"`
IncludeSummary bool `json:"includeSummary"`
IncludeRawContent bool `json:"includeRawContent"`
ConciseSnippet bool `json:"conciseSnippet"`
}
type metasoSearchResponse struct {
Webpages []metasoWebpage `json:"webpages"`
}
type metasoWebpage struct {
Title string `json:"title"`
Link string `json:"link"`
Snippet string `json:"snippet"`
Summary string `json:"summary"`
RawContent string `json:"rawContent"`
Date string `json:"date"`
}
@@ -0,0 +1,87 @@
package web_search
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
)
func TestMetasoProviderSearch(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Fatalf("method = %s, want POST", r.Method)
}
if got := r.Header.Get("Authorization"); got != "Bearer mk-test" {
t.Fatalf("Authorization = %q", got)
}
var request metasoSearchRequest
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Fatal(err)
}
if request.Query != "WeKnora" || request.Scope != "scholar" || request.Size != 2 {
t.Fatalf("unexpected request: %+v", request)
}
if !request.IncludeSummary || request.IncludeRawContent || !request.ConciseSnippet {
t.Fatalf("unexpected content options: %+v", request)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"webpages":[
{"title":"First","link":"https://example.com/1","summary":"Summary","snippet":"Snippet","date":"2026-08-09"},
{"title":"Second","link":"https://example.com/2","snippet":"Fallback snippet","date":"invalid"},
{"title":"Third","link":"https://example.com/3","snippet":"must be capped"}
]}`))
}))
defer server.Close()
metaso := &MetasoProvider{client: server.Client(), baseURL: server.URL, apiKey: "mk-test", scope: "scholar"}
results, err := metaso.Search(context.Background(), " WeKnora ", 2, true)
if err != nil {
t.Fatal(err)
}
if len(results) != 2 {
t.Fatalf("len(results) = %d, want 2", len(results))
}
if results[0].Snippet != "Summary" || results[1].Snippet != "Fallback snippet" {
t.Fatalf("unexpected snippets: %q, %q", results[0].Snippet, results[1].Snippet)
}
if results[0].Source != "metaso" || results[0].PublishedAt == nil {
t.Fatalf("unexpected first result: %+v", results[0])
}
if want := time.Date(2026, 8, 9, 0, 0, 0, 0, time.UTC); !results[0].PublishedAt.Equal(want) {
t.Fatalf("date = %v, want %v", results[0].PublishedAt, want)
}
if results[1].PublishedAt != nil {
t.Fatalf("invalid date should be ignored: %v", results[1].PublishedAt)
}
}
func TestValidateMetasoParameters(t *testing.T) {
if err := ValidateMetasoParameters(types.WebSearchProviderParameters{}); err == nil {
t.Fatal("expected missing API key error")
}
if err := ValidateMetasoParameters(types.WebSearchProviderParameters{APIKey: "mk-test", ExtraConfig: map[string]string{"scope": "unknown"}}); err == nil {
t.Fatal("expected invalid scope error")
}
if err := ValidateMetasoParameters(types.WebSearchProviderParameters{APIKey: "mk-test"}); err != nil {
t.Fatalf("default parameters: %v", err)
}
}
func TestMetasoProviderHTTPError(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"message":"invalid API key"}`))
}))
defer server.Close()
metaso := &MetasoProvider{client: server.Client(), baseURL: server.URL, apiKey: "bad", scope: defaultMetasoScope}
_, err := metaso.Search(context.Background(), "test", 1, false)
if err == nil || !strings.Contains(err.Error(), "invalid API key") {
t.Fatalf("error = %v", err)
}
}
@@ -0,0 +1,156 @@
package web_search
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const (
// defaultOllamaWebSearchURL is the hardcoded Ollama web search API URL.
// Not configurable by tenants — prevents SSRF.
defaultOllamaWebSearchURL = "https://ollama.com/api/web_search"
defaultOllamaTimeout = 10 * time.Second
defaultOllamaResults = 5
maxOllamaResults = 10 // Ollama限制最多10个结果
)
// OllamaProvider implements web search using Ollama Cloud API
type OllamaProvider struct {
client *http.Client
baseURL string
apiKey string
}
// NewOllamaProvider creates a new Ollama web search provider from parameters.
func NewOllamaProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
if params.APIKey == "" {
return nil, fmt.Errorf("API key is required for Ollama provider")
}
client, err := NewSearchHTTPClient(defaultOllamaTimeout, params.ProxyURL)
if err != nil {
return nil, fmt.Errorf("create Ollama HTTP client: %w", err)
}
return &OllamaProvider{
client: client,
baseURL: defaultOllamaWebSearchURL, // Hardcoded — not tenant-configurable
apiKey: params.APIKey,
}, nil
}
// Name returns the provider name
func (p *OllamaProvider) Name() string {
return "ollama"
}
// Search performs a web search using Ollama Cloud API
func (p *OllamaProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
if len(query) == 0 {
return nil, fmt.Errorf("query is empty")
}
if maxResults <= 0 {
maxResults = defaultOllamaResults
}
// Ollama限制最多10个结果
if maxResults > maxOllamaResults {
maxResults = maxOllamaResults
}
logger.Infof(ctx, "[WebSearch][Ollama] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
req, err := p.buildRequest(ctx, query, maxResults)
if err != nil {
return nil, err
}
results, err := p.doSearch(ctx, req)
if err != nil {
logger.Warnf(ctx, "[WebSearch][Ollama] failed: %v", err)
return nil, err
}
logger.Infof(ctx, "[WebSearch][Ollama] returned %d results", len(results))
return results, nil
}
func (p *OllamaProvider) buildRequest(ctx context.Context, query string, maxResults int) (*http.Request, error) {
requestBody := map[string]interface{}{
"query": query,
"max_results": maxResults,
}
bodyBytes, err := json.Marshal(requestBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request body: %w", err)
}
req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL, bytes.NewReader(bodyBytes))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+p.apiKey)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", defaultUserAgentHeader)
return req, nil
}
func (p *OllamaProvider) doSearch(ctx context.Context, req *http.Request) ([]*types.WebSearchResult, error) {
resp, err := p.client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
logger.Warnf(ctx, "[WebSearch][Ollama] API returned status %d: %s", resp.StatusCode, string(body))
return nil, fmt.Errorf("ollama API returned status %d: %s", resp.StatusCode, string(body))
}
var respData ollamaSearchResponse
if err := json.Unmarshal(body, &respData); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
results := make([]*types.WebSearchResult, 0, len(respData.Results))
for _, item := range respData.Results {
results = append(results, &types.WebSearchResult{
Title: item.Title,
URL: item.URL,
Snippet: item.Snippet,
Content: item.Content,
Source: "ollama",
})
}
return results, nil
}
// ollamaSearchResponse defines the response structure for Ollama web search API.
type ollamaSearchResponse struct {
Results []ollamaSearchResult `json:"results"`
}
type ollamaSearchResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
Snippet string `json:"snippet"`
}
@@ -0,0 +1,53 @@
package web_search
import (
"fmt"
"net/http"
"net/url"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/utils"
)
// ValidateProxyURL delegates to utils.ValidateURLForSSRF (only http/https pass that check).
func ValidateProxyURL(proxyURL string) error {
proxyURL = strings.TrimSpace(proxyURL)
if proxyURL == "" {
return nil
}
return utils.ValidateURLForSSRF(proxyURL)
}
// NewSearchHTTPClient builds an http.Client for outbound web search requests.
// It uses utils.SSRFSafeDialContext, optional explicit or environment proxy, and
// redirect validation consistent with utils.NewSSRFSafeHTTPClient.
func NewSearchHTTPClient(timeout time.Duration, proxyURL string) (*http.Client, error) {
proxyURL = strings.TrimSpace(proxyURL)
def, ok := http.DefaultTransport.(*http.Transport)
if !ok {
return nil, fmt.Errorf("default HTTP transport is not *http.Transport")
}
t := def.Clone()
t.DialContext = utils.SSRFSafeDialContext
if proxyURL != "" {
if err := ValidateProxyURL(proxyURL); err != nil {
return nil, err
}
u, err := url.Parse(proxyURL)
if err != nil {
return nil, fmt.Errorf("invalid proxy_url: %w", err)
}
if u.Scheme == "" || u.Host == "" {
return nil, fmt.Errorf("invalid proxy_url: scheme and host are required")
}
t.Proxy = http.ProxyURL(u)
} else {
t.Proxy = http.ProxyFromEnvironment
}
cfg := utils.DefaultSSRFSafeHTTPClientConfig()
cfg.Timeout = timeout
return utils.NewSSRFSafeHTTPClientWithTransport(cfg, t), nil
}
@@ -0,0 +1,81 @@
package web_search
import (
"net/http"
"net/url"
"strings"
"testing"
)
func TestValidateProxyURL(t *testing.T) {
tests := []struct {
name string
raw string
wantErr bool
}{
{"empty", "", false},
{"whitespace only", " ", false},
{"no scheme private IP bare host:port", "192.168.0.1:3128", true},
{"localhost http", "http://localhost:3128", true},
{"direct IPv4 http", "http://192.168.1.1:3128", true},
{"socks5 to IP blocked", "socks5://10.0.0.1:1080", true},
{"bad scheme", "ftp://proxy.example:21", true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := ValidateProxyURL(tt.raw)
if tt.wantErr && err == nil {
t.Fatalf("expected error, got nil")
}
if !tt.wantErr && err != nil {
t.Fatalf("unexpected error: %v", err)
}
})
}
}
func TestNewSearchHTTPClient_RejectsUnsafeProxy(t *testing.T) {
_, err := NewSearchHTTPClient(0, "http://127.0.0.1:3128")
if err == nil {
t.Fatal("expected error for disallowed proxy URL (127.0.0.1)")
}
if !strings.Contains(err.Error(), "SSRF") {
t.Fatalf("expected SSRF-related error, got: %v", err)
}
}
func TestNewSearchHTTPClient_AcceptsEmptyProxy(t *testing.T) {
c, err := NewSearchHTTPClient(0, "")
if err != nil {
t.Fatal(err)
}
if c.Transport == nil {
t.Fatal("expected transport")
}
if c.CheckRedirect == nil {
t.Fatal("expected CheckRedirect for SSRF-safe redirects")
}
}
func TestSsrfSafeRedirect_BlocksNonHTTPSScheme(t *testing.T) {
client, err := NewSearchHTTPClient(0, "")
if err != nil {
t.Fatal(err)
}
req := &http.Request{URL: mustParse(t, "ftp://evil.com/")}
// len(via)==0 so first redirect
err = client.CheckRedirect(req, nil)
if err == nil {
t.Fatal("expected error for ftp redirect")
}
}
func mustParse(t *testing.T, s string) *url.URL {
t.Helper()
u, err := url.Parse(s)
if err != nil {
t.Fatal(err)
}
return u
}
@@ -0,0 +1,45 @@
package web_search
import (
"fmt"
"sync"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
// ProviderFactory creates a new web search provider instance from parameters.
type ProviderFactory func(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error)
// Registry manages web search provider type registrations.
// It maps provider type IDs (e.g., "bing", "google") to their factory functions.
// Instances are created on-demand with tenant-specific parameters.
type Registry struct {
factories map[string]ProviderFactory
mu sync.RWMutex
}
// NewRegistry creates a new web search provider registry
func NewRegistry() *Registry {
return &Registry{
factories: make(map[string]ProviderFactory),
}
}
// Register registers a provider type factory by ID
func (r *Registry) Register(id string, factory ProviderFactory) {
r.mu.Lock()
defer r.mu.Unlock()
r.factories[id] = factory
}
// CreateProvider creates a provider instance by type with the given parameters.
func (r *Registry) CreateProvider(providerType string, params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
r.mu.RLock()
factory, ok := r.factories[providerType]
r.mu.RUnlock()
if !ok {
return nil, fmt.Errorf("web search provider type %s not registered", providerType)
}
return factory(params)
}
@@ -0,0 +1,209 @@
package web_search
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/Tencent/WeKnora/internal/utils"
)
// defaultSearxngTimeout is sized slightly above the SearXNG image's default
// outgoing.max_request_timeout (10s in docker/searxng/settings.yml) so a slow
// upstream engine surfaces as a SearXNG-side error instead of a client cancel.
const defaultSearxngTimeout = 12 * time.Second
// SearxngProvider implements web search using a self-hosted SearXNG instance.
//
// Unlike commercial providers, SearXNG is self-hosted, so the instance URL is
// supplied by the tenant via WebSearchProviderParameters.BaseURL. The URL is
// validated with utils.ValidateURLForSSRF; private/loopback hosts must be added
// to the SSRF_WHITELIST environment variable.
type SearxngProvider struct {
client *http.Client
baseURL string
lastUnresponsive [][]string
}
// ValidateSearxngBaseURL validates a SearXNG instance URL: must be a non-empty,
// absolute http(s) URL, and must pass the SSRF whitelist check. Shared between
// the service-layer parameter validation and the provider constructor so that
// "save" and "use" never disagree.
func ValidateSearxngBaseURL(rawURL string) error {
base := strings.TrimSpace(rawURL)
if base == "" {
return fmt.Errorf("base_url is required for SearXNG provider")
}
parsed, err := url.Parse(base)
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
return fmt.Errorf("invalid SearXNG base_url: must be an absolute http(s) URL")
}
if parsed.Scheme != "http" && parsed.Scheme != "https" {
return fmt.Errorf("invalid SearXNG base_url scheme: %s", parsed.Scheme)
}
if parsed.RawQuery != "" || parsed.Fragment != "" {
return fmt.Errorf("invalid SearXNG base_url: must not contain query or fragment")
}
if err := utils.ValidateURLForSSRF(base); err != nil {
return fmt.Errorf("invalid SearXNG base_url: %w", err)
}
return nil
}
// NewSearxngProvider builds a SearXNG provider from tenant parameters.
func NewSearxngProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
base := strings.TrimSpace(params.BaseURL)
if err := ValidateSearxngBaseURL(base); err != nil {
return nil, err
}
client, err := NewSearchHTTPClient(defaultSearxngTimeout, params.ProxyURL)
if err != nil {
return nil, err
}
return &SearxngProvider{
client: client,
baseURL: strings.TrimRight(base, "/"),
}, nil
}
// Name returns the provider name.
func (p *SearxngProvider) Name() string { return "searxng" }
// EmptyResultDiagnostics explains why the most recent search returned no
// usable results. Used by the settings "test connection" flow.
func (p *SearxngProvider) EmptyResultDiagnostics() string {
if detail := formatUnresponsiveEngines(p.lastUnresponsive); detail != "" {
return detail + "; check that upstream search engines can reach the internet"
}
return "verify the instance URL is reachable and JSON format is enabled in settings.yml"
}
// Search performs a metasearch query against the configured SearXNG instance.
// SearXNG must have `search.formats: [json]` enabled in settings.yml.
func (p *SearxngProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
if strings.TrimSpace(query) == "" {
return nil, fmt.Errorf("query is empty")
}
if maxResults <= 0 {
maxResults = 5
}
q := url.Values{}
q.Set("q", query)
q.Set("format", "json")
// Use "all" (SearXNG's documented value for "no language filter") instead
// of "auto", which is a UI-side default and not a valid /search parameter.
// safesearch is intentionally not set here so the value configured in the
// instance's settings.yml is honored.
q.Set("language", "all")
reqURL := p.baseURL + "/search?" + q.Encode()
logger.Infof(ctx, "[WebSearch][SearXNG] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Accept", "application/json")
req.Header.Set("User-Agent", "WeKnora/1.0")
resp, err := p.client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
return nil, fmt.Errorf("searxng returned status %d: %s", resp.StatusCode, string(body))
}
var data searxngResponse
if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
p.lastUnresponsive = nil
return nil, fmt.Errorf("failed to decode SearXNG response (ensure JSON format is enabled in settings.yml): %w", err)
}
p.lastUnresponsive = data.UnresponsiveEngines
results := make([]*types.WebSearchResult, 0, maxResults)
for _, r := range data.Results {
if len(results) >= maxResults {
break
}
if r.URL == "" || r.Title == "" {
continue
}
item := &types.WebSearchResult{
Title: r.Title,
URL: r.URL,
Snippet: r.Content,
Source: "searxng",
}
if includeDate && r.PublishedDate != "" {
if t, ok := parseSearxngDate(r.PublishedDate); ok {
item.PublishedAt = &t
} else {
logger.Debugf(ctx, "[WebSearch][SearXNG] unparsable publishedDate=%q", r.PublishedDate)
}
}
results = append(results, item)
}
if len(results) == 0 && len(data.UnresponsiveEngines) > 0 {
logger.Warnf(ctx, "[WebSearch][SearXNG] empty results, unresponsive_engines=%v", data.UnresponsiveEngines)
}
logger.Infof(ctx, "[WebSearch][SearXNG] returned %d results", len(results))
return results, nil
}
// searxngDateLayouts covers the formats different SearXNG engines emit for
// publishedDate. Order matters only for performance; first match wins.
// time.RFC3339 already accepts the nanosecond-precision form, so RFC3339Nano
// is intentionally omitted.
var searxngDateLayouts = []string{
time.RFC3339,
"2006-01-02T15:04:05",
"2006-01-02 15:04:05",
"2006-01-02",
time.RFC1123Z,
time.RFC1123,
}
func parseSearxngDate(s string) (time.Time, bool) {
s = strings.TrimSpace(s)
if s == "" {
return time.Time{}, false
}
for _, layout := range searxngDateLayouts {
if t, err := time.Parse(layout, s); err == nil {
return t, true
}
}
return time.Time{}, false
}
type searxngResponse struct {
Results []struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
PublishedDate string `json:"publishedDate,omitempty"`
} `json:"results"`
// UnresponsiveEngines is a list of [engine, reason] tuples returned by
// SearXNG when one or more upstream engines fail. We surface it on empty
// result sets to aid debugging.
UnresponsiveEngines [][]string `json:"unresponsive_engines,omitempty"`
}
@@ -0,0 +1,167 @@
package web_search
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/utils"
)
func TestValidateSearxngBaseURL(t *testing.T) {
// utils.ValidateURLForSSRF caches the parsed SSRF_WHITELIST via
// sync.Once on first call. An alphabetically-earlier test in this
// binary (TestValidateProxyURL) triggers ValidateURLForSSRF with an
// empty whitelist and caches an empty config, so the later setenv
// here would otherwise be ignored. Reset the singleton on both
// entry and exit to keep the env / singleton in sync.
utils.ResetSSRFWhitelistForTest()
os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
defer func() {
os.Unsetenv("SSRF_WHITELIST")
utils.ResetSSRFWhitelistForTest()
}()
cases := []struct {
name string
url string
wantErr bool
}{
{name: "empty", url: "", wantErr: true},
{name: "no scheme", url: "searxng:8080", wantErr: true},
{name: "bad scheme", url: "ftp://searxng:8080", wantErr: true},
{name: "with query", url: "http://127.0.0.1:8080/?x=1", wantErr: true},
{name: "with fragment", url: "http://127.0.0.1:8080/#frag", wantErr: true},
{name: "loopback ok via whitelist", url: "http://127.0.0.1:8888", wantErr: false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
err := ValidateSearxngBaseURL(tc.url)
if (err != nil) != tc.wantErr {
t.Fatalf("ValidateSearxngBaseURL(%q) err=%v wantErr=%v", tc.url, err, tc.wantErr)
}
})
}
}
func TestParseSearxngDate(t *testing.T) {
cases := []struct {
in string
ok bool
}{
{"", false},
{"2024-05-01", true},
{"2024-05-01T12:30:45", true},
{"2024-05-01T12:30:45Z", true},
{"2024-05-01T12:30:45.123456789Z", true},
{"2024-05-01 12:30:45", true},
{"Wed, 01 May 2024 12:30:45 GMT", true},
{"not-a-date", false},
}
for _, tc := range cases {
t.Run(tc.in, func(t *testing.T) {
_, ok := parseSearxngDate(tc.in)
if ok != tc.ok {
t.Fatalf("parseSearxngDate(%q) ok=%v want=%v", tc.in, ok, tc.ok)
}
})
}
}
func TestSearxngProvider_Search(t *testing.T) {
// See TestValidateSearxngBaseURL comment — reset the SSRF whitelist
// singleton so the setenv below is actually observed by the cached
// ssrfWhitelistConfig in internal/utils.
utils.ResetSSRFWhitelistForTest()
os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
defer func() {
os.Unsetenv("SSRF_WHITELIST")
utils.ResetSSRFWhitelistForTest()
}()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/search" {
http.NotFound(w, r)
return
}
if got := r.URL.Query().Get("format"); got != "json" {
t.Errorf("expected format=json, got %q", got)
}
if got := r.URL.Query().Get("q"); got != "hello" {
t.Errorf("expected q=hello, got %q", got)
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"results": []map[string]any{
{"title": "T1", "url": "https://e/1", "content": "c1", "publishedDate": "2024-05-01"},
{"title": "", "url": "https://e/skip"},
{"title": "T2", "url": "https://e/2", "content": "c2"},
},
})
}))
defer srv.Close()
provider, err := NewSearxngProvider(types.WebSearchProviderParameters{BaseURL: srv.URL})
if err != nil {
t.Fatalf("NewSearxngProvider: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
results, err := provider.Search(ctx, "hello", 5, true)
if err != nil {
t.Fatalf("Search: %v", err)
}
if len(results) != 2 {
t.Fatalf("expected 2 results (one skipped), got %d", len(results))
}
if results[0].PublishedAt == nil {
t.Fatalf("expected first result PublishedAt to be set")
}
if got := results[0].Source; got != "searxng" {
t.Fatalf("unexpected source: %q", got)
}
}
func TestSearxngProvider_Search_EmptyWithUnresponsiveEngines(t *testing.T) {
utils.ResetSSRFWhitelistForTest()
os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
defer func() {
os.Unsetenv("SSRF_WHITELIST")
utils.ResetSSRFWhitelistForTest()
}()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"results": []any{},
"unresponsive_engines": [][]string{{"google", "timeout"}},
})
}))
defer srv.Close()
provider, err := NewSearxngProvider(types.WebSearchProviderParameters{BaseURL: srv.URL})
if err != nil {
t.Fatalf("NewSearxngProvider: %v", err)
}
sp := provider.(*SearxngProvider)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
results, err := sp.Search(ctx, "test", 1, false)
if err != nil {
t.Fatalf("Search: %v", err)
}
if len(results) != 0 {
t.Fatalf("expected 0 results, got %d", len(results))
}
if got := sp.EmptyResultDiagnostics(); !strings.Contains(got, "google (timeout)") {
t.Fatalf("EmptyResultDiagnostics() = %q", got)
}
}
@@ -0,0 +1,142 @@
package web_search
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const (
// defaultTavilySearchURL is the hardcoded Tavily API URL.
// Not configurable by tenants — prevents SSRF.
defaultTavilySearchURL = "https://api.tavily.com/search"
)
var (
defaultTavilyTimeout = 15 * time.Second
)
// TavilyProvider implements web search using Tavily Search API
type TavilyProvider struct {
client *http.Client
baseURL string
apiKey string
}
// NewTavilyProvider creates a new Tavily provider from parameters (no environment variables).
func NewTavilyProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
if params.APIKey == "" {
return nil, fmt.Errorf("API key is required for Tavily provider")
}
client, err := NewSearchHTTPClient(defaultTavilyTimeout, params.ProxyURL)
if err != nil {
return nil, err
}
return &TavilyProvider{
client: client,
baseURL: defaultTavilySearchURL,
apiKey: params.APIKey,
}, nil
}
// Name returns the provider name
func (p *TavilyProvider) Name() string {
return "tavily"
}
// Search performs a web search using Tavily Search API
func (p *TavilyProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
if len(query) == 0 {
return nil, fmt.Errorf("query is empty")
}
logger.Infof(ctx, "[WebSearch][Tavily] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
reqBody := tavilySearchRequest{
APIKey: p.apiKey,
Query: query,
MaxResults: maxResults,
}
bodyBytes, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL, bytes.NewReader(bodyBytes))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
resp, err := p.client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
respBody, _ := io.ReadAll(resp.Body)
logger.Warnf(ctx, "[WebSearch][Tavily] API returned status %d: %s", resp.StatusCode, string(respBody))
return nil, fmt.Errorf("tavily API returned status %d: %s", resp.StatusCode, string(respBody))
}
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
var respData tavilySearchResponse
if err := json.Unmarshal(respBody, &respData); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
results := make([]*types.WebSearchResult, 0, len(respData.Results))
for _, item := range respData.Results {
result := &types.WebSearchResult{
Title: item.Title,
URL: item.URL,
Snippet: item.Content,
Source: "tavily",
}
if includeDate && item.PublishedDate != "" {
if t, err := time.Parse(time.RFC3339, item.PublishedDate); err == nil {
result.PublishedAt = &t
}
}
results = append(results, result)
}
logger.Infof(ctx, "[WebSearch][Tavily] returned %d results", len(results))
return results, nil
}
// tavilySearchRequest defines the request body for Tavily search API
type tavilySearchRequest struct {
APIKey string `json:"api_key"`
Query string `json:"query"`
MaxResults int `json:"max_results"`
}
// tavilySearchResponse defines the response structure for Tavily search API
type tavilySearchResponse struct {
Query string `json:"query"`
Results []struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
Score float64 `json:"score"`
PublishedDate string `json:"published_date,omitempty"`
} `json:"results"`
}
@@ -0,0 +1,66 @@
package web_search
import (
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/types"
)
// emptyResultDiagnostics is implemented by providers that can explain an empty
// search result set (used by the settings "test connection" flow).
type emptyResultDiagnostics interface {
EmptyResultDiagnostics() string
}
// EmptyTestResultsError builds a provider-specific message when a connectivity
// test search succeeds but returns no usable results.
func EmptyTestResultsError(providerType string, provider any) error {
detail := ""
if dr, ok := provider.(emptyResultDiagnostics); ok {
detail = dr.EmptyResultDiagnostics()
}
switch types.WebSearchProviderType(providerType) {
case types.WebSearchProviderTypeSearxng:
if detail != "" {
return fmt.Errorf("searxng returned 0 results: %s", detail)
}
return fmt.Errorf(
"searxng returned 0 results; verify the instance URL, JSON format in settings.yml, and upstream search engine connectivity",
)
case types.WebSearchProviderTypeDuckDuckGo:
return fmt.Errorf("duckduckgo returned 0 results; verify network connectivity and proxy settings")
case types.WebSearchProviderTypeKeenable:
return fmt.Errorf(
"keenable returned 0 results; keyless requests are rate-limited, so verify network connectivity or set an API key to lift the cap",
)
case types.WebSearchProviderTypeExa:
return fmt.Errorf(
"exa returned 0 results; verify the API key, account quota, network connectivity, and proxy settings",
)
default:
return fmt.Errorf("search returned 0 results, please verify your API key and configuration")
}
}
func formatUnresponsiveEngines(engines [][]string) string {
if len(engines) == 0 {
return ""
}
parts := make([]string, 0, len(engines))
for _, e := range engines {
if len(e) == 0 {
continue
}
if len(e) == 1 {
parts = append(parts, e[0])
continue
}
parts = append(parts, e[0]+" ("+e[1]+")")
}
if len(parts) == 0 {
return ""
}
return "unresponsive engines: " + strings.Join(parts, ", ")
}
@@ -0,0 +1,102 @@
package web_search
import (
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestEmptyTestResultsError_SearxngWithUnresponsiveEngines(t *testing.T) {
provider := &SearxngProvider{
lastUnresponsive: [][]string{
{"google", "timeout"},
{"bing", "blocked"},
},
}
err := EmptyTestResultsError(string(types.WebSearchProviderTypeSearxng), provider)
if err == nil {
t.Fatal("expected error")
}
msg := err.Error()
if !strings.Contains(msg, "searxng returned 0 results") {
t.Fatalf("unexpected message: %q", msg)
}
if !strings.Contains(msg, "google (timeout)") || !strings.Contains(msg, "bing (blocked)") {
t.Fatalf("expected unresponsive engine details, got: %q", msg)
}
}
func TestEmptyTestResultsError_SearxngWithoutDiagnostics(t *testing.T) {
provider := &SearxngProvider{}
err := EmptyTestResultsError(string(types.WebSearchProviderTypeSearxng), provider)
if err == nil {
t.Fatal("expected error")
}
msg := err.Error()
if !strings.Contains(msg, "JSON format is enabled in settings.yml") {
t.Fatalf("unexpected message: %q", msg)
}
if strings.Contains(msg, "API key") {
t.Fatalf("searxng message must not mention API key: %q", msg)
}
}
func TestEmptyTestResultsError_DuckDuckGo(t *testing.T) {
err := EmptyTestResultsError(string(types.WebSearchProviderTypeDuckDuckGo), nil)
if err == nil {
t.Fatal("expected error")
}
if !strings.Contains(err.Error(), "duckduckgo returned 0 results") {
t.Fatalf("unexpected message: %q", err.Error())
}
}
func TestEmptyTestResultsError_Keenable(t *testing.T) {
err := EmptyTestResultsError(string(types.WebSearchProviderTypeKeenable), nil)
if err == nil {
t.Fatal("expected error")
}
msg := err.Error()
if !strings.Contains(msg, "keenable returned 0 results") {
t.Fatalf("unexpected message: %q", msg)
}
if !strings.Contains(msg, "rate-limited") {
t.Fatalf("keenable message should hint at the keyless rate limit: %q", msg)
}
}
func TestEmptyTestResultsError_Exa(t *testing.T) {
err := EmptyTestResultsError(string(types.WebSearchProviderTypeExa), nil)
if err == nil {
t.Fatal("expected error")
}
msg := err.Error()
if !strings.Contains(msg, "exa returned 0 results") {
t.Fatalf("unexpected message: %q", msg)
}
if !strings.Contains(msg, "API key") || !strings.Contains(msg, "quota") || !strings.Contains(msg, "proxy") {
t.Fatalf("exa message should include actionable connection checks: %q", msg)
}
}
func TestEmptyTestResultsError_APIKeyProviders(t *testing.T) {
err := EmptyTestResultsError(string(types.WebSearchProviderTypeBing), nil)
if err == nil {
t.Fatal("expected error")
}
if !strings.Contains(err.Error(), "API key") {
t.Fatalf("bing message should mention API key: %q", err.Error())
}
}
func TestFormatUnresponsiveEngines(t *testing.T) {
got := formatUnresponsiveEngines([][]string{
{"google", "timeout"},
{"wikipedia"},
})
want := "unresponsive engines: google (timeout), wikipedia"
if got != want {
t.Fatalf("formatUnresponsiveEngines() = %q, want %q", got, want)
}
}
+271
View File
@@ -0,0 +1,271 @@
package web_search
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
const (
defaultZhipuSearchURL = "https://open.bigmodel.cn/api/paas/v4/web_search"
defaultZhipuTimeout = 15 * time.Second
defaultZhipuResults = 10
maxZhipuResults = 50
maxZhipuQueryRunes = 70
maxZhipuResponseBytes = 2 << 20
defaultZhipuSearchEngine = "search_std"
defaultZhipuContentSize = "medium"
)
var validZhipuSearchEngines = map[string]struct{}{
"search_std": {},
"search_pro": {},
"search_pro_sogou": {},
"search_pro_quark": {},
}
var validZhipuContentSizes = map[string]struct{}{
"medium": {},
"high": {},
}
// ZhipuProvider implements web search using Zhipu AI's standalone Web Search API.
type ZhipuProvider struct {
client *http.Client
baseURL string
apiKey string
searchEngine string
contentSize string
}
// NewZhipuProvider creates a Zhipu AI web search provider from persisted parameters.
func NewZhipuProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
if err := ValidateZhipuParameters(params); err != nil {
return nil, err
}
client, err := NewSearchHTTPClient(defaultZhipuTimeout, params.ProxyURL)
if err != nil {
return nil, err
}
searchEngine, contentSize := zhipuOptions(params.ExtraConfig)
return &ZhipuProvider{
client: client,
baseURL: defaultZhipuSearchURL,
apiKey: strings.TrimSpace(params.APIKey),
searchEngine: searchEngine,
contentSize: contentSize,
}, nil
}
// ValidateZhipuParameters validates credentials and provider-specific options.
func ValidateZhipuParameters(params types.WebSearchProviderParameters) error {
if strings.TrimSpace(params.APIKey) == "" {
return fmt.Errorf("API key is required for Zhipu provider")
}
searchEngine, contentSize := zhipuOptions(params.ExtraConfig)
if _, ok := validZhipuSearchEngines[searchEngine]; !ok {
return fmt.Errorf("invalid Zhipu search engine: %s", searchEngine)
}
if _, ok := validZhipuContentSizes[contentSize]; !ok {
return fmt.Errorf("invalid Zhipu content size: %s", contentSize)
}
return nil
}
func zhipuOptions(extraConfig map[string]string) (searchEngine, contentSize string) {
searchEngine = defaultZhipuSearchEngine
contentSize = defaultZhipuContentSize
if value := strings.TrimSpace(extraConfig["search_engine"]); value != "" {
searchEngine = value
}
if value := strings.TrimSpace(extraConfig["content_size"]); value != "" {
contentSize = value
}
return searchEngine, contentSize
}
// Name returns the provider name.
func (p *ZhipuProvider) Name() string {
return "zhipu"
}
// Search performs a web search using Zhipu AI's standalone Web Search API.
func (p *ZhipuProvider) Search(
ctx context.Context,
query string,
maxResults int,
includeDate bool,
) ([]*types.WebSearchResult, error) {
preparedQuery := normalizeZhipuQuery(query)
if preparedQuery == "" {
return nil, fmt.Errorf("query is empty")
}
if utf8.RuneCountInString(strings.TrimSpace(query)) > maxZhipuQueryRunes {
logger.Infof(ctx, "[WebSearch][Zhipu] truncated query to %d characters", maxZhipuQueryRunes)
}
if maxResults <= 0 {
maxResults = defaultZhipuResults
}
if maxResults > maxZhipuResults {
maxResults = maxZhipuResults
}
requestBody := zhipuSearchRequest{
SearchQuery: preparedQuery,
SearchEngine: p.searchEngine,
SearchIntent: false,
Count: maxResults,
ContentSize: p.contentSize,
}
body, err := json.Marshal(requestBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal Zhipu request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("failed to create Zhipu request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+p.apiKey)
req.Header.Set("Content-Type", "application/json")
logger.Infof(ctx, "[WebSearch][Zhipu] query=%q maxResults=%d engine=%s", preparedQuery, maxResults, p.searchEngine)
resp, err := p.client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute Zhipu request: %w", err)
}
defer resp.Body.Close()
respBody, err := readZhipuResponseBody(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
return nil, zhipuHTTPError(resp.StatusCode, respBody)
}
var response zhipuSearchResponse
if err := json.Unmarshal(respBody, &response); err != nil {
return nil, fmt.Errorf("failed to unmarshal Zhipu response: %w", err)
}
if response.Error.Message != "" || response.Error.Code != "" {
return nil, fmt.Errorf("Zhipu API error (%s): %s", response.Error.Code, response.Error.Message)
}
results := make([]*types.WebSearchResult, 0, len(response.SearchResult))
for _, item := range response.SearchResult {
if strings.TrimSpace(item.Title) == "" && strings.TrimSpace(item.Link) == "" {
continue
}
result := &types.WebSearchResult{
Title: item.Title,
URL: item.Link,
Snippet: item.Content,
Source: "zhipu",
}
if includeDate {
if publishedAt, ok := parseZhipuDate(item.PublishDate); ok {
result.PublishedAt = &publishedAt
}
}
results = append(results, result)
if len(results) >= maxResults {
break
}
}
logger.Infof(ctx, "[WebSearch][Zhipu] returned %d results", len(results))
return results, nil
}
func normalizeZhipuQuery(query string) string {
query = strings.TrimSpace(query)
if utf8.RuneCountInString(query) <= maxZhipuQueryRunes {
return query
}
runes := []rune(query)
return string(runes[:maxZhipuQueryRunes])
}
func parseZhipuDate(value string) (time.Time, bool) {
value = strings.TrimSpace(value)
if value == "" {
return time.Time{}, false
}
for _, layout := range []string{
time.RFC3339Nano,
"2006-01-02 15:04:05",
"2006-01-02 15:04",
"2006-01-02",
} {
if parsed, err := time.Parse(layout, value); err == nil {
return parsed, true
}
}
return time.Time{}, false
}
func readZhipuResponseBody(reader io.Reader) ([]byte, error) {
body, err := io.ReadAll(io.LimitReader(reader, maxZhipuResponseBytes+1))
if err != nil {
return nil, fmt.Errorf("failed to read Zhipu response: %w", err)
}
if len(body) > maxZhipuResponseBytes {
return nil, fmt.Errorf("Zhipu response exceeds %d bytes", maxZhipuResponseBytes)
}
return body, nil
}
func zhipuHTTPError(statusCode int, body []byte) error {
var response zhipuSearchResponse
if err := json.Unmarshal(body, &response); err == nil && (response.Error.Code != "" || response.Error.Message != "") {
return fmt.Errorf("Zhipu API returned status %d (%s): %s", statusCode, response.Error.Code, response.Error.Message)
}
detail := strings.TrimSpace(string(body))
if len(detail) > 4096 {
detail = detail[:4096]
}
if detail == "" {
return fmt.Errorf("Zhipu API returned status %d", statusCode)
}
return fmt.Errorf("Zhipu API returned status %d: %s", statusCode, detail)
}
type zhipuSearchRequest struct {
SearchQuery string `json:"search_query"`
SearchEngine string `json:"search_engine"`
SearchIntent bool `json:"search_intent"`
Count int `json:"count"`
ContentSize string `json:"content_size"`
}
type zhipuSearchResponse struct {
ID string `json:"id"`
RequestID string `json:"request_id"`
SearchResult []zhipuSearchResult `json:"search_result"`
Error zhipuError `json:"error"`
}
type zhipuSearchResult struct {
Title string `json:"title"`
Content string `json:"content"`
Link string `json:"link"`
Media string `json:"media"`
Icon string `json:"icon"`
Refer string `json:"refer"`
PublishDate string `json:"publish_date"`
}
type zhipuError struct {
Code string `json:"code"`
Message string `json:"message"`
}
@@ -0,0 +1,202 @@
package web_search
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/types"
)
func TestValidateZhipuParameters(t *testing.T) {
tests := []struct {
name string
params types.WebSearchProviderParameters
wantErr bool
}{
{name: "defaults", params: types.WebSearchProviderParameters{APIKey: "key"}},
{
name: "custom options",
params: types.WebSearchProviderParameters{
APIKey: "key",
ExtraConfig: map[string]string{
"search_engine": "search_pro_sogou",
"content_size": "high",
},
},
},
{name: "missing key", params: types.WebSearchProviderParameters{}, wantErr: true},
{
name: "invalid engine",
params: types.WebSearchProviderParameters{
APIKey: "key",
ExtraConfig: map[string]string{"search_engine": "unknown"},
},
wantErr: true,
},
{
name: "invalid content size",
params: types.WebSearchProviderParameters{
APIKey: "key",
ExtraConfig: map[string]string{"content_size": "large"},
},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := ValidateZhipuParameters(tt.params)
if (err != nil) != tt.wantErr {
t.Fatalf("ValidateZhipuParameters() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}
func TestZhipuProviderSearch(t *testing.T) {
query := strings.Repeat("智", maxZhipuQueryRunes+1)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Errorf("method = %s, want POST", r.Method)
}
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
t.Errorf("Authorization = %q, want Bearer test-key", got)
}
var request zhipuSearchRequest
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Errorf("decode request: %v", err)
}
if got := utf8.RuneCountInString(request.SearchQuery); got != maxZhipuQueryRunes {
t.Errorf("query rune count = %d, want %d", got, maxZhipuQueryRunes)
}
if request.SearchEngine != "search_pro" {
t.Errorf("search_engine = %q, want search_pro", request.SearchEngine)
}
if request.ContentSize != "high" {
t.Errorf("content_size = %q, want high", request.ContentSize)
}
if request.SearchIntent {
t.Error("search_intent = true, want false")
}
if request.Count != maxZhipuResults {
t.Errorf("count = %d, want %d", request.Count, maxZhipuResults)
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "search-id",
"request_id": "request-id",
"search_result": []map[string]any{
{
"title": "Result 1",
"link": "https://example.com/1",
"content": "Summary 1",
"publish_date": "2026-07-16",
},
{
"title": "Result 2",
"link": "https://example.com/2",
"content": "Summary 2",
},
},
})
}))
defer srv.Close()
provider := &ZhipuProvider{
client: srv.Client(),
baseURL: srv.URL,
apiKey: "test-key",
searchEngine: "search_pro",
contentSize: "high",
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
results, err := provider.Search(ctx, query, maxZhipuResults+1, true)
if err != nil {
t.Fatalf("Search() error = %v", err)
}
if len(results) != 2 {
t.Fatalf("len(results) = %d, want 2", len(results))
}
if results[0].Snippet != "Summary 1" || results[0].Content != "" {
t.Errorf("first result content mapping = %+v", results[0])
}
if results[0].Source != "zhipu" {
t.Errorf("source = %q, want zhipu", results[0].Source)
}
if results[0].PublishedAt == nil || results[0].PublishedAt.Format("2006-01-02") != "2026-07-16" {
t.Errorf("published_at = %v, want 2026-07-16", results[0].PublishedAt)
}
}
func TestZhipuProviderSearchDefaults(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var request zhipuSearchRequest
_ = json.NewDecoder(r.Body).Decode(&request)
if request.SearchEngine != defaultZhipuSearchEngine {
t.Errorf("search_engine = %q, want %q", request.SearchEngine, defaultZhipuSearchEngine)
}
if request.ContentSize != defaultZhipuContentSize {
t.Errorf("content_size = %q, want %q", request.ContentSize, defaultZhipuContentSize)
}
if request.Count != defaultZhipuResults {
t.Errorf("count = %d, want %d", request.Count, defaultZhipuResults)
}
_, _ = w.Write([]byte(`{"search_result":[]}`))
}))
defer srv.Close()
provider := &ZhipuProvider{
client: srv.Client(),
baseURL: srv.URL,
apiKey: "test-key",
searchEngine: defaultZhipuSearchEngine,
contentSize: defaultZhipuContentSize,
}
if _, err := provider.Search(context.Background(), "test", 0, false); err != nil {
t.Fatalf("Search() error = %v", err)
}
}
func TestZhipuProviderSearchError(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusTooManyRequests)
_, _ = w.Write([]byte(`{"error":{"code":"1302","message":"rate limited"}}`))
}))
defer srv.Close()
provider := &ZhipuProvider{
client: srv.Client(),
baseURL: srv.URL,
apiKey: "test-key",
searchEngine: defaultZhipuSearchEngine,
contentSize: defaultZhipuContentSize,
}
_, err := provider.Search(context.Background(), "test", 1, false)
if err == nil || !strings.Contains(err.Error(), "rate limited") {
t.Fatalf("Search() error = %v, want rate limited error", err)
}
}
func TestParseZhipuDate(t *testing.T) {
for _, value := range []string{
"2026-07-16",
"2026-07-16 12:30",
"2026-07-16 12:30:45",
"2026-07-16T12:30:45Z",
} {
if _, ok := parseZhipuDate(value); !ok {
t.Errorf("parseZhipuDate(%q) failed", value)
}
}
if _, ok := parseZhipuDate("not-a-date"); ok {
t.Error("parseZhipuDate(not-a-date) unexpectedly succeeded")
}
}