初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,125 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
const braveAPIBase = "https://api.search.brave.com"
|
||||
|
||||
// BraveProvider 实现 SearchProvider。
|
||||
type BraveProvider struct {
|
||||
client *http.Client
|
||||
apiBase string
|
||||
maxResults int
|
||||
}
|
||||
|
||||
func NewBraveProvider(timeout time.Duration, maxResults int) *BraveProvider {
|
||||
return &BraveProvider{
|
||||
client: &http.Client{
|
||||
Timeout: timeout,
|
||||
},
|
||||
apiBase: braveAPIBase,
|
||||
maxResults: maxResults,
|
||||
}
|
||||
}
|
||||
|
||||
func (b *BraveProvider) Name() models.ProviderKind {
|
||||
return models.ProviderBrave
|
||||
}
|
||||
|
||||
func (b *BraveProvider) MaxResults() int {
|
||||
return b.maxResults
|
||||
}
|
||||
|
||||
type braveWebResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Description string `json:"description"`
|
||||
Age string `json:"age"`
|
||||
}
|
||||
|
||||
type braveResponse struct {
|
||||
Web struct {
|
||||
Results []braveWebResult `json:"results"`
|
||||
} `json:"web"`
|
||||
}
|
||||
|
||||
func (b *BraveProvider) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
|
||||
maxRes := req.MaxResults
|
||||
if maxRes <= 0 || maxRes > b.maxResults {
|
||||
maxRes = b.maxResults
|
||||
}
|
||||
|
||||
u, _ := url.Parse(b.apiBase + "/res/v1/web/search")
|
||||
q := u.Query()
|
||||
q.Set("q", req.Query)
|
||||
q.Set("count", fmt.Sprintf("%d", maxRes))
|
||||
u.RawQuery = q.Encode()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("brave: create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Accept", "application/json")
|
||||
httpReq.Header.Set("X-Subscription-Token", apiKey)
|
||||
|
||||
start := time.Now()
|
||||
resp, err := b.client.Do(httpReq)
|
||||
duration := time.Since(start)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("brave: do request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("brave: read body: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("brave: status %d: %s", resp.StatusCode, string(raw))
|
||||
}
|
||||
|
||||
var br braveResponse
|
||||
if err := json.Unmarshal(raw, &br); err != nil {
|
||||
return nil, fmt.Errorf("brave: decode: %w", err)
|
||||
}
|
||||
|
||||
items := make([]models.ResultItem, 0, len(br.Web.Results))
|
||||
for _, r := range br.Web.Results {
|
||||
items = append(items, models.ResultItem{
|
||||
Title: r.Title,
|
||||
URL: r.URL,
|
||||
Snippet: r.Description,
|
||||
PublishedAt: r.Age,
|
||||
})
|
||||
}
|
||||
|
||||
return &models.SearchResult{
|
||||
Provider: models.ProviderBrave,
|
||||
Query: req.Query,
|
||||
Results: items,
|
||||
TotalCount: len(items),
|
||||
DurationMs: duration.Milliseconds(),
|
||||
Cached: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (b *BraveProvider) HealthCheck(ctx context.Context) error {
|
||||
var d net.Dialer
|
||||
conn, err := d.DialContext(ctx, "tcp", "api.search.brave.com:443")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
conn.Close()
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
func TestBraveProvider_Search_Success(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("X-Subscription-Token") == "" {
|
||||
t.Error("missing subscription token header")
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(`{
|
||||
"web": {
|
||||
"results": [
|
||||
{"title": "Brave Result", "url": "https://brave.com/1", "description": "Brave snippet"}
|
||||
]
|
||||
}
|
||||
}`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
p := &BraveProvider{
|
||||
client: ts.Client(),
|
||||
apiBase: ts.URL,
|
||||
maxResults: 10,
|
||||
}
|
||||
|
||||
r, err := p.Search(context.Background(), "test-key", &models.SearchRequest{Query: "test"})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if r.TotalCount != 1 {
|
||||
t.Fatalf("expected 1 result, got %d", r.TotalCount)
|
||||
}
|
||||
if r.Results[0].Title != "Brave Result" {
|
||||
t.Errorf("expected 'Brave Result', got '%s'", r.Results[0].Title)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBraveProvider_Name(t *testing.T) {
|
||||
p := NewBraveProvider(5*time.Second, 10)
|
||||
if p.Name() != models.ProviderBrave {
|
||||
t.Errorf("expected brave, got %s", p.Name())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
// Package providers 封装搜索提供商适配。
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// SearchProvider 搜索提供商接口。
|
||||
type SearchProvider interface {
|
||||
// Name 返回提供商标识。
|
||||
Name() models.ProviderKind
|
||||
// Search 执行搜索。apiKey 由负载均衡器传入，provider 本身不管理密钥。
|
||||
Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error)
|
||||
// HealthCheck 健康检查。
|
||||
HealthCheck(ctx context.Context) error
|
||||
// MaxResults 返回单次搜索最大数量。
|
||||
MaxResults() int
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// Registry 管理多 provider 注册和查询。
|
||||
type Registry struct {
|
||||
mu sync.RWMutex
|
||||
providers map[models.ProviderKind]SearchProvider
|
||||
}
|
||||
|
||||
func NewRegistry() *Registry {
|
||||
return &Registry{
|
||||
providers: make(map[models.ProviderKind]SearchProvider),
|
||||
}
|
||||
}
|
||||
|
||||
// Register 注册 provider。
|
||||
func (r *Registry) Register(p SearchProvider) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.providers[p.Name()] = p
|
||||
}
|
||||
|
||||
// Get 按名称获取 provider。
|
||||
func (r *Registry) Get(name models.ProviderKind) (SearchProvider, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
p, ok := r.providers[name]
|
||||
return p, ok
|
||||
}
|
||||
|
||||
// All 返回所有已注册 provider。
|
||||
func (r *Registry) All() []SearchProvider {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
out := make([]SearchProvider, 0, len(r.providers))
|
||||
for _, p := range r.providers {
|
||||
out = append(out, p)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Names 返回所有已注册的 provider 名称。
|
||||
func (r *Registry) Names() []models.ProviderKind {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
out := make([]models.ProviderKind, 0, len(r.providers))
|
||||
for name := range r.providers {
|
||||
out = append(out, name)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// SetupFromConfig 根据配置创建并注册所有启用的 provider。
|
||||
func SetupFromConfig(cfg *config.Config) (*Registry, error) {
|
||||
reg := NewRegistry()
|
||||
|
||||
if cfg.Providers.Tavily.Enabled {
|
||||
if len(cfg.Providers.Tavily.APIKeys) == 0 {
|
||||
slog.Warn("tavily: enabled but no API keys configured, skipping")
|
||||
} else {
|
||||
reg.Register(NewTavilyProvider(cfg.Providers.Tavily.Timeout, cfg.Providers.Tavily.MaxResults))
|
||||
}
|
||||
}
|
||||
if cfg.Providers.Brave.Enabled {
|
||||
if len(cfg.Providers.Brave.APIKeys) == 0 {
|
||||
slog.Warn("brave: enabled but no API keys configured, skipping")
|
||||
} else {
|
||||
reg.Register(NewBraveProvider(cfg.Providers.Brave.Timeout, cfg.Providers.Brave.MaxResults))
|
||||
}
|
||||
}
|
||||
|
||||
if len(reg.Names()) == 0 {
|
||||
return nil, fmt.Errorf("no providers enabled or configured")
|
||||
}
|
||||
return reg, nil
|
||||
}
|
||||
|
||||
// SetupFromConfigWithOverrides 从配置创建注册中心，并允许使用显式 provider 配置覆盖配置文件。
|
||||
// 本函数方便测试和动态配置场景。
|
||||
func SetupFromConfigWithOverrides(
|
||||
cfg *config.Config,
|
||||
overrideTimeout time.Duration,
|
||||
overrideMaxResults int,
|
||||
) (*Registry, error) {
|
||||
reg := NewRegistry()
|
||||
|
||||
if cfg.Providers.Tavily.Enabled && len(cfg.Providers.Tavily.APIKeys) > 0 {
|
||||
t := cfg.Providers.Tavily.Timeout
|
||||
m := cfg.Providers.Tavily.MaxResults
|
||||
if overrideTimeout > 0 {
|
||||
t = overrideTimeout
|
||||
}
|
||||
if overrideMaxResults > 0 {
|
||||
m = overrideMaxResults
|
||||
}
|
||||
reg.Register(NewTavilyProvider(t, m))
|
||||
}
|
||||
if cfg.Providers.Brave.Enabled && len(cfg.Providers.Brave.APIKeys) > 0 {
|
||||
t := cfg.Providers.Brave.Timeout
|
||||
m := cfg.Providers.Brave.MaxResults
|
||||
if overrideTimeout > 0 {
|
||||
t = overrideTimeout
|
||||
}
|
||||
if overrideMaxResults > 0 {
|
||||
m = overrideMaxResults
|
||||
}
|
||||
reg.Register(NewBraveProvider(t, m))
|
||||
}
|
||||
|
||||
if len(reg.Names()) == 0 {
|
||||
return nil, fmt.Errorf("no providers enabled or configured")
|
||||
}
|
||||
return reg, nil
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
type mockProvider struct {
|
||||
name models.ProviderKind
|
||||
}
|
||||
|
||||
func (m *mockProvider) Name() models.ProviderKind { return m.name }
|
||||
func (m *mockProvider) Search(ctx context.Context, s string, req *models.SearchRequest) (*models.SearchResult, error) {
|
||||
return &models.SearchResult{Provider: m.name, Query: req.Query}, nil
|
||||
}
|
||||
func (m *mockProvider) HealthCheck(ctx context.Context) error { return nil }
|
||||
func (m *mockProvider) MaxResults() int { return 10 }
|
||||
|
||||
func TestRegistry_RegisterAndGet(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
m := &mockProvider{name: models.ProviderTavily}
|
||||
r.Register(m)
|
||||
|
||||
got, ok := r.Get(models.ProviderTavily)
|
||||
if !ok {
|
||||
t.Fatal("expected to find provider")
|
||||
}
|
||||
if got.Name() != models.ProviderTavily {
|
||||
t.Errorf("expected tavily, got %s", got.Name())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_All(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
r.Register(&mockProvider{name: models.ProviderTavily})
|
||||
r.Register(&mockProvider{name: models.ProviderBrave})
|
||||
|
||||
all := r.All()
|
||||
if len(all) != 2 {
|
||||
t.Errorf("expected 2 providers, got %d", len(all))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_GetMissing(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
_, ok := r.Get(models.ProviderBrave)
|
||||
if ok {
|
||||
t.Fatal("expected missing provider")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistry_Names(t *testing.T) {
|
||||
r := NewRegistry()
|
||||
r.Register(&mockProvider{name: models.ProviderTavily})
|
||||
r.Register(&mockProvider{name: models.ProviderBrave})
|
||||
|
||||
names := r.Names()
|
||||
if len(names) != 2 {
|
||||
t.Errorf("expected 2 names, got %d", len(names))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
const tavilyAPIBase = "https://api.tavily.com"
|
||||
|
||||
type tavilyRequest struct {
|
||||
APIKey string `json:"api_key"`
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
type tavilyResponse struct {
|
||||
Answer string `json:"answer"`
|
||||
Query string `json:"query"`
|
||||
ResponseTime float64 `json:"response_time"`
|
||||
Results []tavilyResult `json:"results"`
|
||||
}
|
||||
|
||||
type tavilyResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
Score float64 `json:"score"`
|
||||
PublishedAt string `json:"published_date"`
|
||||
}
|
||||
|
||||
// TavilyProvider 实现 SearchProvider。
|
||||
type TavilyProvider struct {
|
||||
client *http.Client
|
||||
apiBase string
|
||||
maxResults int
|
||||
}
|
||||
|
||||
func NewTavilyProvider(timeout time.Duration, maxResults int) *TavilyProvider {
|
||||
return &TavilyProvider{
|
||||
client: &http.Client{
|
||||
Timeout: timeout,
|
||||
},
|
||||
apiBase: tavilyAPIBase,
|
||||
maxResults: maxResults,
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TavilyProvider) Name() models.ProviderKind {
|
||||
return models.ProviderTavily
|
||||
}
|
||||
|
||||
func (t *TavilyProvider) MaxResults() int {
|
||||
return t.maxResults
|
||||
}
|
||||
|
||||
func (t *TavilyProvider) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
|
||||
maxRes := req.MaxResults
|
||||
if maxRes <= 0 || maxRes > t.maxResults {
|
||||
maxRes = t.maxResults
|
||||
}
|
||||
|
||||
body := tavilyRequest{
|
||||
APIKey: apiKey,
|
||||
Query: req.Query,
|
||||
MaxResults: maxRes,
|
||||
}
|
||||
payload, _ := json.Marshal(body)
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, t.apiBase+"/search", bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("tavily: create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
start := time.Now()
|
||||
resp, err := t.client.Do(httpReq)
|
||||
duration := time.Since(start)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("tavily: do request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("tavily: read body: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("tavily: status %d: %s", resp.StatusCode, string(raw))
|
||||
}
|
||||
|
||||
var tr tavilyResponse
|
||||
if err := json.Unmarshal(raw, &tr); err != nil {
|
||||
return nil, fmt.Errorf("tavily: decode response: %w", err)
|
||||
}
|
||||
|
||||
items := make([]models.ResultItem, 0, len(tr.Results))
|
||||
for _, r := range tr.Results {
|
||||
items = append(items, models.ResultItem{
|
||||
Title: r.Title,
|
||||
URL: r.URL,
|
||||
Snippet: r.Content,
|
||||
Score: r.Score,
|
||||
PublishedAt: r.PublishedAt,
|
||||
})
|
||||
}
|
||||
|
||||
return &models.SearchResult{
|
||||
Provider: models.ProviderTavily,
|
||||
Query: req.Query,
|
||||
Results: items,
|
||||
TotalCount: len(items),
|
||||
DurationMs: duration.Milliseconds(),
|
||||
Cached: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (t *TavilyProvider) HealthCheck(ctx context.Context) error {
|
||||
var d net.Dialer
|
||||
conn, err := d.DialContext(ctx, "tcp", "api.tavily.com:443")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
conn.Close()
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
func TestTavilyProvider_Search_Success(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/search" {
|
||||
t.Errorf("unexpected path: %s", r.URL.Path)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(`{
|
||||
"query": "test",
|
||||
"results": [
|
||||
{"title": "Result 1", "url": "https://example.com/1", "content": "Content 1", "score": 0.95}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
p := &TavilyProvider{
|
||||
client: ts.Client(),
|
||||
apiBase: ts.URL,
|
||||
maxResults: 10,
|
||||
}
|
||||
|
||||
result, err := p.Search(context.Background(), "test-key", &models.SearchRequest{Query: "test"})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result.TotalCount != 1 {
|
||||
t.Fatalf("expected 1 result, got %d", result.TotalCount)
|
||||
}
|
||||
if result.Results[0].Title != "Result 1" {
|
||||
t.Errorf("expected 'Result 1', got '%s'", result.Results[0].Title)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTavilyProvider_Search_HTTPError(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
w.Write([]byte(`{"error": "invalid api key"}`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
p := &TavilyProvider{
|
||||
client: ts.Client(),
|
||||
apiBase: ts.URL,
|
||||
maxResults: 10,
|
||||
}
|
||||
|
||||
_, err := p.Search(context.Background(), "bad-key", &models.SearchRequest{Query: "test"})
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTavilyProvider_Name(t *testing.T) {
|
||||
p := NewTavilyProvider(5*time.Second, 10)
|
||||
if p.Name() != models.ProviderTavily {
|
||||
t.Errorf("expected tavily, got %s", p.Name())
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user