初始化提交
This commit is contained in:
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,264 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// KeySearcher 按密钥执行搜索的接口。
|
||||
// 与 SearchProvider 解耦,方便测试时注入 mock。
|
||||
type KeySearcher interface {
|
||||
Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error)
|
||||
HealthCheck(ctx context.Context) error
|
||||
Name() models.ProviderKind
|
||||
}
|
||||
|
||||
// LoadBalancer 负载均衡器(goroutine-safe)。
|
||||
type LoadBalancer struct {
|
||||
pools map[models.ProviderKind]*KeyPool
|
||||
searchers map[models.ProviderKind]KeySearcher
|
||||
disabled map[models.ProviderKind]bool // 运行时动态禁用
|
||||
retryCount int
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewLoadBalancer(cfg *config.BalancerConfig) *LoadBalancer {
|
||||
return &LoadBalancer{
|
||||
pools: make(map[models.ProviderKind]*KeyPool),
|
||||
searchers: make(map[models.ProviderKind]KeySearcher),
|
||||
disabled: make(map[models.ProviderKind]bool),
|
||||
retryCount: cfg.RetryCount,
|
||||
}
|
||||
}
|
||||
|
||||
// AddProvider 注册 provider 及其密钥池。
|
||||
func (lb *LoadBalancer) AddProvider(
|
||||
provider models.ProviderKind,
|
||||
keys []string,
|
||||
weight int,
|
||||
searcher KeySearcher,
|
||||
) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
|
||||
cbThreshold := 5
|
||||
cbCooldown := 30 * time.Second
|
||||
lb.pools[provider] = NewKeyPool(provider, keys, weight, cbThreshold, cbCooldown)
|
||||
lb.searchers[provider] = searcher
|
||||
// 添加时默认启用
|
||||
delete(lb.disabled, provider)
|
||||
}
|
||||
|
||||
// RemoveProvider 移除 provider 及其密钥池。
|
||||
func (lb *LoadBalancer) RemoveProvider(provider models.ProviderKind) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
delete(lb.pools, provider)
|
||||
delete(lb.searchers, provider)
|
||||
delete(lb.disabled, provider)
|
||||
}
|
||||
|
||||
// EnableProvider 启用 provider。
|
||||
func (lb *LoadBalancer) EnableProvider(provider models.ProviderKind) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
delete(lb.disabled, provider)
|
||||
}
|
||||
|
||||
// DisableProvider 禁用 provider(搜索时跳过,不删除配置)。
|
||||
func (lb *LoadBalancer) DisableProvider(provider models.ProviderKind) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
lb.disabled[provider] = true
|
||||
}
|
||||
|
||||
// IsProviderEnabled 检查 provider 是否启用。
|
||||
func (lb *LoadBalancer) IsProviderEnabled(provider models.ProviderKind) bool {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
return !lb.disabled[provider]
|
||||
}
|
||||
|
||||
// SetRetryCount 更新重试次数。
|
||||
func (lb *LoadBalancer) SetRetryCount(n int) {
|
||||
lb.mu.Lock()
|
||||
defer lb.mu.Unlock()
|
||||
lb.retryCount = n
|
||||
}
|
||||
|
||||
// HasProvider 检查 provider 是否已注册。
|
||||
func (lb *LoadBalancer) HasProvider(provider models.ProviderKind) bool {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
_, ok := lb.pools[provider]
|
||||
return ok
|
||||
}
|
||||
|
||||
// Search 执行搜索(自动选密钥、跨 provider 重试)。
|
||||
func (lb *LoadBalancer) Search(ctx context.Context, req *models.SearchRequest) (*models.SearchResult, error) {
|
||||
providers := lb.resolveProviders(req.Providers)
|
||||
|
||||
var lastErr error
|
||||
for attempt := 0; attempt <= lb.getRetryCount(); attempt++ {
|
||||
for _, prov := range providers {
|
||||
pool, ok := lb.pools[prov]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
ks := pool.Select()
|
||||
if ks == nil {
|
||||
continue
|
||||
}
|
||||
searcher := lb.searchers[prov]
|
||||
if searcher == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
result, err := searcher.Search(ctx, ks.Key, req)
|
||||
duration := time.Since(start)
|
||||
|
||||
if err != nil {
|
||||
ks.RecordFailure()
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
ks.RecordSuccess(duration)
|
||||
result.Provider = prov
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
|
||||
if lastErr != nil {
|
||||
return nil, fmt.Errorf("balancer: all providers exhausted: %w", lastErr)
|
||||
}
|
||||
return nil, fmt.Errorf("balancer: no available provider or key for query: %s", req.Query)
|
||||
}
|
||||
|
||||
// SearchAll 并发搜索所有启用 provider(用于 Merge 模式)。
|
||||
func (lb *LoadBalancer) SearchAll(ctx context.Context, req *models.SearchRequest) []*models.SearchResult {
|
||||
providers := lb.resolveProviders(nil)
|
||||
type result struct {
|
||||
prov models.ProviderKind
|
||||
res *models.SearchResult
|
||||
err error
|
||||
}
|
||||
|
||||
ch := make(chan result, len(providers))
|
||||
for _, prov := range providers {
|
||||
p := prov
|
||||
go func() {
|
||||
pool, ok := lb.pools[p]
|
||||
if !ok {
|
||||
ch <- result{prov: p, err: fmt.Errorf("no pool")}
|
||||
return
|
||||
}
|
||||
ks := pool.Select()
|
||||
if ks == nil {
|
||||
ch <- result{prov: p, err: fmt.Errorf("no available key")}
|
||||
return
|
||||
}
|
||||
searcher := lb.searchers[p]
|
||||
if searcher == nil {
|
||||
ch <- result{prov: p, err: fmt.Errorf("no searcher")}
|
||||
return
|
||||
}
|
||||
start := time.Now()
|
||||
res, err := searcher.Search(ctx, ks.Key, req)
|
||||
duration := time.Since(start)
|
||||
if err != nil {
|
||||
ks.RecordFailure()
|
||||
ch <- result{prov: p, err: err}
|
||||
return
|
||||
}
|
||||
ks.RecordSuccess(duration)
|
||||
res.Provider = p
|
||||
ch <- result{prov: p, res: res}
|
||||
}()
|
||||
}
|
||||
|
||||
var results []*models.SearchResult
|
||||
for range providers {
|
||||
r := <-ch
|
||||
if r.res != nil {
|
||||
results = append(results, r.res)
|
||||
}
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
|
||||
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
|
||||
// ActiveKeyCount 返回每个 provider 的活跃密钥数。
|
||||
func (lb *LoadBalancer) ActiveKeyCount() map[models.ProviderKind]int {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
out := make(map[models.ProviderKind]int)
|
||||
for p, pool := range lb.pools {
|
||||
if lb.disabled[p] {
|
||||
continue
|
||||
}
|
||||
out[p] = pool.ActiveCount()
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Providers 返回所有已注册的 provider 名称(无论启用与否)。
|
||||
func (lb *LoadBalancer) Providers() []models.ProviderKind {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
out := make([]models.ProviderKind, 0, len(lb.pools))
|
||||
for p := range lb.pools {
|
||||
out = append(out, p)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// HealthCheck 对指定 provider 执行健康检查(委托给 KeySearcher)。
|
||||
func (lb *LoadBalancer) HealthCheck(ctx context.Context, provider models.ProviderKind) error {
|
||||
lb.mu.RLock()
|
||||
searcher, ok := lb.searchers[provider]
|
||||
lb.mu.RUnlock()
|
||||
if !ok {
|
||||
return fmt.Errorf("healthcheck: provider not found: %s", provider)
|
||||
}
|
||||
return searcher.HealthCheck(ctx)
|
||||
}
|
||||
|
||||
func (lb *LoadBalancer) getRetryCount() int {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
return lb.retryCount
|
||||
}
|
||||
|
||||
// resolveProviders 按用户偏好和可用性确定 provider 顺序(跳过已禁用的)。
|
||||
func (lb *LoadBalancer) resolveProviders(preferred []string) []models.ProviderKind {
|
||||
lb.mu.RLock()
|
||||
defer lb.mu.RUnlock()
|
||||
|
||||
if len(preferred) == 0 {
|
||||
out := make([]models.ProviderKind, 0, len(lb.pools))
|
||||
for p := range lb.pools {
|
||||
if !lb.disabled[p] {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
out := make([]models.ProviderKind, 0, len(preferred))
|
||||
seen := make(map[models.ProviderKind]bool)
|
||||
for _, name := range preferred {
|
||||
p := models.ProviderKind(name)
|
||||
if _, ok := lb.pools[p]; ok && !seen[p] && !lb.disabled[p] {
|
||||
out = append(out, p)
|
||||
seen[p] = true
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
type mockSearcher struct {
|
||||
name models.ProviderKind
|
||||
failOn int // 第几次调用失败(0 = 永不)
|
||||
calls int
|
||||
}
|
||||
|
||||
func (m *mockSearcher) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
|
||||
m.calls++
|
||||
if m.failOn > 0 && m.calls >= m.failOn {
|
||||
return nil, errors.New("mock failure")
|
||||
}
|
||||
return &models.SearchResult{Provider: m.name, Query: req.Query, Results: []models.ResultItem{{Title: "mock"}}}, nil
|
||||
}
|
||||
func (m *mockSearcher) HealthCheck(ctx context.Context) error { return nil }
|
||||
func (m *mockSearcher) Name() models.ProviderKind { return m.name }
|
||||
|
||||
func TestLoadBalancer_Search_Basic(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
|
||||
|
||||
result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result.Provider != models.ProviderTavily {
|
||||
t.Errorf("expected tavily, got %s", result.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBalancer_Search_FallbackProvider(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
|
||||
lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})
|
||||
|
||||
result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if string(result.Provider) != "brave" {
|
||||
t.Errorf("expected fallback to brave, got %s", result.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBalancer_Search_AllExhausted(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
|
||||
|
||||
_, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
|
||||
if err == nil {
|
||||
t.Fatal("expected error when all keys exhausted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBalancer_Search_PreferredProviders(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
|
||||
lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})
|
||||
|
||||
result, err := lb.Search(context.Background(), &models.SearchRequest{
|
||||
Query: "test",
|
||||
Providers: []string{"brave"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result.Provider != models.ProviderBrave {
|
||||
t.Errorf("expected bing (preferred), got %s", result.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBalancer_ActiveKeyCount(t *testing.T) {
|
||||
lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
|
||||
lb.AddProvider(models.ProviderTavily, []string{"key1", "key2"}, 1, &mockSearcher{name: models.ProviderTavily})
|
||||
|
||||
counts := lb.ActiveKeyCount()
|
||||
if counts[models.ProviderTavily] != 2 {
|
||||
t.Errorf("expected 2 active keys, got %d", counts[models.ProviderTavily])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// HealthChecker 定期对已注册 provider 执行健康检查。
|
||||
// 连续失败超过阈值时自动禁用 provider,恢复成功后自动启用。
|
||||
type HealthChecker struct {
|
||||
lb *LoadBalancer
|
||||
interval time.Duration
|
||||
timeout time.Duration
|
||||
threshold int
|
||||
|
||||
mu sync.Mutex
|
||||
failures map[models.ProviderKind]int // 连续失败计数
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// NewHealthChecker 创建健康检查器。cfg 为对应 provider 的健康检查配置。
|
||||
func NewHealthChecker(lb *LoadBalancer, cfg config.HealthCheckConf) *HealthChecker {
|
||||
return &HealthChecker{
|
||||
lb: lb,
|
||||
interval: cfg.Interval,
|
||||
timeout: cfg.Timeout,
|
||||
threshold: cfg.FailureThreshold,
|
||||
failures: make(map[models.ProviderKind]int),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start 在后台 goroutine 中启动健康检查循环。
|
||||
// 每隔 interval 时间对所有已注册 provider 执行一次健康检查。
|
||||
func (hc *HealthChecker) Start() {
|
||||
if hc.interval <= 0 {
|
||||
return
|
||||
}
|
||||
slog.Info("健康检查已启动", "interval", hc.interval, "timeout", hc.timeout, "failure_threshold", hc.threshold)
|
||||
go hc.loop()
|
||||
}
|
||||
|
||||
// Stop 停止健康检查循环。
|
||||
func (hc *HealthChecker) Stop() {
|
||||
close(hc.done)
|
||||
}
|
||||
|
||||
func (hc *HealthChecker) loop() {
|
||||
ticker := time.NewTicker(hc.interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
hc.checkAll()
|
||||
case <-hc.done:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (hc *HealthChecker) checkAll() {
|
||||
// 获取当前所有已注册 provider 的名称
|
||||
// 我们需要从 LoadBalancer 获取提供者列表
|
||||
// 由于 LoadBalancer 没有直接暴露 provider 列表的 API,我们遍历 ActiveKeyCount
|
||||
// 来获取所有已知的 provider
|
||||
for _, prov := range hc.getAllProviders() {
|
||||
hc.checkProvider(prov)
|
||||
}
|
||||
}
|
||||
|
||||
func (hc *HealthChecker) checkProvider(prov models.ProviderKind) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), hc.timeout)
|
||||
defer cancel()
|
||||
|
||||
err := hc.lb.HealthCheck(ctx, prov)
|
||||
|
||||
hc.mu.Lock()
|
||||
defer hc.mu.Unlock()
|
||||
|
||||
if err != nil {
|
||||
hc.failures[prov]++
|
||||
fails := hc.failures[prov]
|
||||
slog.Warn("健康检查失败",
|
||||
"provider", prov,
|
||||
"consecutive_failures", fails,
|
||||
"error", err,
|
||||
)
|
||||
if fails >= hc.threshold && hc.lb.IsProviderEnabled(prov) {
|
||||
hc.lb.DisableProvider(prov)
|
||||
slog.Warn("健康检查: provider 已自动禁用",
|
||||
"provider", prov,
|
||||
"failures", fails,
|
||||
)
|
||||
}
|
||||
} else {
|
||||
// 恢复成功
|
||||
if hc.failures[prov] > 0 {
|
||||
slog.Info("健康检查恢复", "provider", prov)
|
||||
}
|
||||
hc.failures[prov] = 0
|
||||
if !hc.lb.IsProviderEnabled(prov) {
|
||||
hc.lb.EnableProvider(prov)
|
||||
slog.Info("健康检查: provider 已自动启用", "provider", prov)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getAllProviders 返回当前所有已注册的 provider 名称(包括已禁用的)。
|
||||
func (hc *HealthChecker) getAllProviders() []models.ProviderKind {
|
||||
return hc.lb.Providers()
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
// Package balancer 实现密钥池管理和负载均衡。
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// State 熔断器状态。
|
||||
type State int
|
||||
|
||||
const (
|
||||
StateClosed State = iota // 正常
|
||||
StateOpen // 熔断
|
||||
StateHalfOpen // 半开(探测)
|
||||
)
|
||||
|
||||
func (s State) String() string {
|
||||
switch s {
|
||||
case StateClosed:
|
||||
return "closed"
|
||||
case StateOpen:
|
||||
return "open"
|
||||
case StateHalfOpen:
|
||||
return "half-open"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// CircuitBreaker 熔断器。
|
||||
type CircuitBreaker struct {
|
||||
mu sync.Mutex
|
||||
state State
|
||||
failures int
|
||||
threshold int
|
||||
cooldown time.Duration
|
||||
lastFailureTime time.Time
|
||||
}
|
||||
|
||||
func NewCircuitBreaker(threshold int, cooldown time.Duration) *CircuitBreaker {
|
||||
return &CircuitBreaker{
|
||||
state: StateClosed,
|
||||
threshold: threshold,
|
||||
cooldown: cooldown,
|
||||
}
|
||||
}
|
||||
|
||||
// Allow 判断是否允许请求通过。
|
||||
func (cb *CircuitBreaker) Allow() bool {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
|
||||
switch cb.state {
|
||||
case StateClosed:
|
||||
return true
|
||||
case StateOpen:
|
||||
if time.Since(cb.lastFailureTime) > cb.cooldown {
|
||||
cb.state = StateHalfOpen
|
||||
return true
|
||||
}
|
||||
return false
|
||||
case StateHalfOpen:
|
||||
return true
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// Success 记录成功。
|
||||
func (cb *CircuitBreaker) Success() {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
cb.failures = 0
|
||||
if cb.state == StateHalfOpen {
|
||||
cb.state = StateClosed
|
||||
}
|
||||
}
|
||||
|
||||
// Failure 记录失败。
|
||||
func (cb *CircuitBreaker) Failure() {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
cb.failures++
|
||||
cb.lastFailureTime = time.Now()
|
||||
if cb.state == StateHalfOpen || cb.failures >= cb.threshold {
|
||||
cb.state = StateOpen
|
||||
}
|
||||
}
|
||||
|
||||
// State 返回当前状态(goroutine-safe 快照)。
|
||||
func (cb *CircuitBreaker) State() State {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
return cb.state
|
||||
}
|
||||
|
||||
// Failures 返回连续失败次数。
|
||||
func (cb *CircuitBreaker) Failures() int {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
return cb.failures
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// KeyStatus 密钥状态。
|
||||
type KeyStatus int
|
||||
|
||||
const (
|
||||
KeyActive KeyStatus = iota // 正常
|
||||
KeyCooldown // 冷却中(熔断后)
|
||||
KeyDead // 死亡(多次失败不再恢复)
|
||||
)
|
||||
|
||||
// KeyState 单个密钥状态。
|
||||
type KeyState struct {
|
||||
Provider models.ProviderKind
|
||||
Key string
|
||||
Weight int
|
||||
|
||||
mu sync.Mutex
|
||||
status KeyStatus
|
||||
latencyEWMA time.Duration // 指数加权滑动平均延迟
|
||||
successCount int64
|
||||
failCount int64
|
||||
cb *CircuitBreaker
|
||||
}
|
||||
|
||||
func newKeyState(provider models.ProviderKind, key string, weight int, cbThreshold int, cbCooldown time.Duration) *KeyState {
|
||||
return &KeyState{
|
||||
Provider: provider,
|
||||
Key: key,
|
||||
Weight: weight,
|
||||
status: KeyActive,
|
||||
cb: NewCircuitBreaker(cbThreshold, cbCooldown),
|
||||
}
|
||||
}
|
||||
|
||||
// Allow 检查是否允许使用此密钥。
|
||||
func (ks *KeyState) Allow() bool {
|
||||
ks.mu.Lock()
|
||||
defer ks.mu.Unlock()
|
||||
if ks.status != KeyActive {
|
||||
return false
|
||||
}
|
||||
return ks.cb.Allow()
|
||||
}
|
||||
|
||||
// RecordSuccess 记录成功调用,更新 EWMA 延迟。
|
||||
func (ks *KeyState) RecordSuccess(latency time.Duration) {
|
||||
ks.mu.Lock()
|
||||
defer ks.mu.Unlock()
|
||||
ks.successCount++
|
||||
ks.cb.Success()
|
||||
// EWMA 更新:α = 0.3
|
||||
alpha := 0.3
|
||||
if ks.latencyEWMA == 0 {
|
||||
ks.latencyEWMA = latency
|
||||
} else {
|
||||
ks.latencyEWMA = time.Duration(
|
||||
(1-alpha)*float64(ks.latencyEWMA) + alpha*float64(latency),
|
||||
)
|
||||
}
|
||||
ks.status = KeyActive
|
||||
}
|
||||
|
||||
// RecordFailure 记录失败调用。
|
||||
func (ks *KeyState) RecordFailure() {
|
||||
ks.mu.Lock()
|
||||
defer ks.mu.Unlock()
|
||||
ks.failCount++
|
||||
ks.cb.Failure()
|
||||
if ks.cb.State() == StateOpen {
|
||||
ks.status = KeyCooldown
|
||||
}
|
||||
}
|
||||
|
||||
// CombinedScore 计算综合分数(越高优先选用)。
|
||||
// 综合 = 基础权重 × 延迟惩罚 × 可靠性因子
|
||||
func (ks *KeyState) CombinedScore() float64 {
|
||||
ks.mu.Lock()
|
||||
defer ks.mu.Unlock()
|
||||
|
||||
if ks.status != KeyActive {
|
||||
return 0
|
||||
}
|
||||
if ks.cb.State() != StateClosed {
|
||||
return 0
|
||||
}
|
||||
|
||||
base := float64(ks.Weight)
|
||||
latencyPenalty := 1.0
|
||||
if ks.latencyEWMA > 0 {
|
||||
latencyPenalty = 1.0 / (1.0 + ks.latencyEWMA.Seconds())
|
||||
}
|
||||
reliability := 1.0
|
||||
total := ks.successCount + ks.failCount
|
||||
if total > 0 && ks.failCount > 0 {
|
||||
reliability = 1.0 - float64(ks.failCount)/float64(total)
|
||||
}
|
||||
return base * latencyPenalty * math.Max(0.1, reliability)
|
||||
}
|
||||
|
||||
// KeyPool 管理一个 provider 下的多密钥。
|
||||
type KeyPool struct {
|
||||
provider models.ProviderKind
|
||||
keys []*KeyState
|
||||
mu sync.RWMutex
|
||||
cbThreshold int
|
||||
cbCooldown time.Duration
|
||||
}
|
||||
|
||||
func NewKeyPool(provider models.ProviderKind, keys []string, weight int, cbThreshold int, cbCooldown time.Duration) *KeyPool {
|
||||
kp := &KeyPool{
|
||||
provider: provider,
|
||||
keys: make([]*KeyState, 0, len(keys)),
|
||||
cbThreshold: cbThreshold,
|
||||
cbCooldown: cbCooldown,
|
||||
}
|
||||
for _, k := range keys {
|
||||
kp.keys = append(kp.keys, newKeyState(provider, k, weight, cbThreshold, cbCooldown))
|
||||
}
|
||||
return kp
|
||||
}
|
||||
|
||||
// Select 按综合分数选最佳密钥。
|
||||
func (kp *KeyPool) Select() *KeyState {
|
||||
kp.mu.RLock()
|
||||
defer kp.mu.RUnlock()
|
||||
|
||||
var best *KeyState
|
||||
var bestScore float64
|
||||
for _, ks := range kp.keys {
|
||||
if !ks.Allow() {
|
||||
continue
|
||||
}
|
||||
score := ks.CombinedScore()
|
||||
if score > bestScore {
|
||||
bestScore = score
|
||||
best = ks
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
// ActiveCount 返回活跃密钥数。
|
||||
func (kp *KeyPool) ActiveCount() int {
|
||||
kp.mu.RLock()
|
||||
defer kp.mu.RUnlock()
|
||||
count := 0
|
||||
for _, ks := range kp.keys {
|
||||
if ks.status == KeyActive {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
// TotalCount 返回密钥总数。
|
||||
func (kp *KeyPool) TotalCount() int {
|
||||
kp.mu.RLock()
|
||||
defer kp.mu.RUnlock()
|
||||
return len(kp.keys)
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
func TestKeyPool_Select_Basic(t *testing.T) {
|
||||
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 5, 30*time.Second)
|
||||
ks := kp.Select()
|
||||
if ks == nil {
|
||||
t.Fatal("expected to select a key")
|
||||
}
|
||||
if ks.Key != "key1" && ks.Key != "key2" {
|
||||
t.Errorf("unexpected key: %s", ks.Key)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyPool_Select_SkipsDead(t *testing.T) {
|
||||
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 2, 1*time.Hour)
|
||||
ks1 := kp.Select()
|
||||
|
||||
// Fail twice to trigger circuit breaker
|
||||
ks1.RecordFailure()
|
||||
ks1.RecordFailure()
|
||||
|
||||
selected := kp.Select()
|
||||
if selected == nil {
|
||||
t.Fatal("expected fallback to second key")
|
||||
}
|
||||
if selected.Key == ks1.Key {
|
||||
t.Error("should not select the failed key")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyState_CombinedScore(t *testing.T) {
|
||||
ks := newKeyState(models.ProviderTavily, "test-key", 2, 5, 30*time.Second)
|
||||
score := ks.CombinedScore()
|
||||
if score <= 0 {
|
||||
t.Errorf("expected positive score, got %f", score)
|
||||
}
|
||||
|
||||
// After recording a slow latency, score should drop
|
||||
ks.RecordSuccess(2 * time.Second)
|
||||
slowScore := ks.CombinedScore()
|
||||
if slowScore >= score {
|
||||
t.Error("expected slow latency to reduce score")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_OpenAndRecover(t *testing.T) {
|
||||
cb := NewCircuitBreaker(3, 50*time.Millisecond)
|
||||
|
||||
// Should be closed initially
|
||||
if !cb.Allow() {
|
||||
t.Error("expected allow initially")
|
||||
}
|
||||
|
||||
// Fail 3 times to open
|
||||
cb.Failure()
|
||||
cb.Failure()
|
||||
cb.Failure()
|
||||
if cb.Allow() {
|
||||
t.Error("expected circuit open after 3 failures")
|
||||
}
|
||||
|
||||
// Wait for cooldown
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
if !cb.Allow() {
|
||||
t.Error("expected allow after cooldown")
|
||||
}
|
||||
|
||||
// Half-open: success should close
|
||||
cb.Success()
|
||||
if cb.State() != StateClosed {
|
||||
t.Error("expected circuit closed after half-open success")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_HalfOpenFailure(t *testing.T) {
|
||||
cb := NewCircuitBreaker(3, 50*time.Millisecond)
|
||||
|
||||
cb.Failure()
|
||||
cb.Failure()
|
||||
cb.Failure()
|
||||
if cb.Allow() {
|
||||
t.Error("expected open")
|
||||
}
|
||||
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
if !cb.Allow() {
|
||||
t.Error("expected half-open allow")
|
||||
}
|
||||
|
||||
// Half-open: failure should re-open immediately
|
||||
cb.Failure()
|
||||
if cb.State() != StateOpen {
|
||||
t.Error("expected open after half-open failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyPool_ActiveCount(t *testing.T) {
|
||||
kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2", "key3"}, 1, 3, 30*time.Second)
|
||||
if kp.TotalCount() != 3 {
|
||||
t.Errorf("expected 3 total, got %d", kp.TotalCount())
|
||||
}
|
||||
if kp.ActiveCount() != 3 {
|
||||
t.Errorf("expected 3 active, got %d", kp.ActiveCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyPool_AllKeysDead(t *testing.T) {
|
||||
kp := NewKeyPool(models.ProviderTavily, []string{"key1"}, 1, 1, 1*time.Hour)
|
||||
ks := kp.Select()
|
||||
ks.RecordFailure()
|
||||
|
||||
selected := kp.Select()
|
||||
if selected != nil {
|
||||
t.Error("expected nil when all keys dead")
|
||||
}
|
||||
}
|
||||
+222
@@ -0,0 +1,222 @@
|
||||
// Package cache 提供搜索缓存(两级:内存 LRU + 可选 Redis)。
|
||||
package cache
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
lru "github.com/hashicorp/golang-lru/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// entry 缓存条目。
|
||||
type entry struct {
|
||||
result *models.SearchResult
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
// Manager 两级缓存管理器。
|
||||
// L1: 内存 LRU(本地毫秒级,单进程共享)
|
||||
// L2: Redis(跨实例共享,可按开关关闭)
|
||||
type Manager struct {
|
||||
enabled bool
|
||||
ttl time.Duration
|
||||
lru *lru.Cache[string, *entry]
|
||||
rdb *redis.Client
|
||||
redisDB int
|
||||
mu sync.RWMutex
|
||||
l1Hits int64
|
||||
l2Hits int64
|
||||
misses int64
|
||||
}
|
||||
|
||||
// New 创建缓存管理器。若 Redis 启用则初始化 L2 连接。
|
||||
func New(cfg config.CacheConfig) (*Manager, error) {
|
||||
if !cfg.Enabled {
|
||||
return &Manager{enabled: false}, nil
|
||||
}
|
||||
|
||||
c, err := lru.New[string, *entry](cfg.MaxItems)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cache: create lru: %w", err)
|
||||
}
|
||||
|
||||
m := &Manager{
|
||||
enabled: true,
|
||||
ttl: cfg.TTL,
|
||||
lru: c,
|
||||
}
|
||||
|
||||
// 初始化 Redis 连接(仅 enabled 时)
|
||||
if cfg.Redis.Enabled {
|
||||
m.rdb = redis.NewClient(&redis.Options{
|
||||
Addr: cfg.Redis.Addr,
|
||||
Password: cfg.Redis.Password,
|
||||
DB: cfg.Redis.DB,
|
||||
})
|
||||
m.redisDB = cfg.Redis.DB
|
||||
// 探测 Redis 连通性
|
||||
if err := m.rdb.Ping(context.Background()).Err(); err != nil {
|
||||
return nil, fmt.Errorf("cache: redis ping failed: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// cacheKey 生成缓存键。格式:bowen:cache:{provider}:{hash[:8]}:{max_results}
|
||||
func cacheKey(req *models.SearchRequest, provider models.ProviderKind) string {
|
||||
h := sha256.Sum256([]byte(fmt.Sprintf("%s:%s:%d", provider, req.Query, req.MaxResults)))
|
||||
return fmt.Sprintf("bowen:cache:%x:%d", h[:8], req.MaxResults)
|
||||
}
|
||||
|
||||
// Get 获取缓存结果。
|
||||
// 查询顺序:L1 (LRU) → L2 (Redis) → nil
|
||||
// 若 L2 命中,自动升温到 L1。
|
||||
func (m *Manager) Get(req *models.SearchRequest, provider models.ProviderKind) *models.SearchResult {
|
||||
if !m.enabled {
|
||||
return nil
|
||||
}
|
||||
key := cacheKey(req, provider)
|
||||
|
||||
// L1: 内存 LRU
|
||||
if result := m.getL1(key); result != nil {
|
||||
m.mu.Lock()
|
||||
m.l1Hits++
|
||||
m.mu.Unlock()
|
||||
return result
|
||||
}
|
||||
|
||||
// L2: Redis(如启用)
|
||||
if m.rdb != nil {
|
||||
if result := m.getL2(key); result != nil {
|
||||
// 升温到 L1
|
||||
m.setL1(key, result)
|
||||
m.mu.Lock()
|
||||
m.l2Hits++
|
||||
m.mu.Unlock()
|
||||
return result
|
||||
}
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
m.misses++
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) getL1(key string) *models.SearchResult {
|
||||
m.mu.RLock()
|
||||
e, ok := m.lru.Get(key)
|
||||
m.mu.RUnlock()
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if time.Now().After(e.expiresAt) {
|
||||
m.mu.Lock()
|
||||
m.lru.Remove(key)
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
result := *e.result
|
||||
result.Cached = true
|
||||
return &result
|
||||
}
|
||||
|
||||
func (m *Manager) getL2(key string) *models.SearchResult {
|
||||
data, err := m.rdb.Get(context.Background(), key).Bytes()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
var result models.SearchResult
|
||||
if err := json.Unmarshal(data, &result); err != nil {
|
||||
return nil
|
||||
}
|
||||
result.Cached = true
|
||||
return &result
|
||||
}
|
||||
|
||||
// Set 写入缓存(L1 + L2,若 Redis 启用则同步写 L2)。
|
||||
func (m *Manager) Set(req *models.SearchRequest, provider models.ProviderKind, result *models.SearchResult) {
|
||||
if !m.enabled {
|
||||
return
|
||||
}
|
||||
key := cacheKey(req, provider)
|
||||
|
||||
// 写 L1
|
||||
m.setL1(key, result)
|
||||
|
||||
// 写 L2(异步 fire-and-forget,不阻塞)
|
||||
if m.rdb != nil {
|
||||
go func() {
|
||||
data, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
m.rdb.Set(context.Background(), key, data, m.ttl)
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) setL1(key string, result *models.SearchResult) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
r := *result
|
||||
m.lru.Add(key, &entry{
|
||||
result: &r,
|
||||
expiresAt: time.Now().Add(m.ttl),
|
||||
})
|
||||
}
|
||||
|
||||
// Stats 缓存统计(L1 命中、L2 命中、L1+L2 总未命中)。
|
||||
func (m *Manager) Stats() (l1Hits, l2Hits, misses int64) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.l1Hits, m.l2Hits, m.misses
|
||||
}
|
||||
|
||||
// HitRate 综合缓存命中率(L1+L2 命中 / 总请求)。
|
||||
func (m *Manager) HitRate() float64 {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
total := m.l1Hits + m.l2Hits + m.misses
|
||||
if total == 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(m.l1Hits+m.l2Hits) / float64(total)
|
||||
}
|
||||
|
||||
// SetTTL 运行时更新 TTL(仅影响后续写入的条目)。
|
||||
func (m *Manager) SetTTL(d time.Duration) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.ttl = d
|
||||
}
|
||||
|
||||
// Clear 清空缓存。
|
||||
func (m *Manager) Clear() {
|
||||
if !m.enabled {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.lru.Purge()
|
||||
m.l1Hits = 0
|
||||
m.l2Hits = 0
|
||||
m.misses = 0
|
||||
}
|
||||
|
||||
// Close 关闭 Redis 连接。
|
||||
func (m *Manager) Close() error {
|
||||
if m.rdb != nil {
|
||||
return m.rdb.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
+122
@@ -0,0 +1,122 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
func TestCache_SetAndGet(t *testing.T) {
|
||||
m, err := New(config.CacheConfig{
|
||||
Enabled: true,
|
||||
TTL: 5 * time.Minute,
|
||||
MaxItems: 100,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
req := &models.SearchRequest{Query: "test", MaxResults: 5}
|
||||
result := &models.SearchResult{Query: "test", TotalCount: 1}
|
||||
|
||||
m.Set(req, models.ProviderTavily, result)
|
||||
got := m.Get(req, models.ProviderTavily)
|
||||
if got == nil {
|
||||
t.Fatal("expected cached result")
|
||||
}
|
||||
if !got.Cached {
|
||||
t.Error("expected cached=true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCache_Miss(t *testing.T) {
|
||||
m, err := New(config.CacheConfig{
|
||||
Enabled: true,
|
||||
TTL: 5 * time.Minute,
|
||||
MaxItems: 100,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
got := m.Get(&models.SearchRequest{Query: "nosuch"}, models.ProviderTavily)
|
||||
if got != nil {
|
||||
t.Fatal("expected nil for cache miss")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCache_Disabled(t *testing.T) {
|
||||
m, err := New(config.CacheConfig{Enabled: false})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
m.Set(&models.SearchRequest{Query: "test"}, models.ProviderTavily, &models.SearchResult{})
|
||||
got := m.Get(&models.SearchRequest{Query: "test"}, models.ProviderTavily)
|
||||
if got != nil {
|
||||
t.Fatal("expected nil when cache disabled")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCache_Expired(t *testing.T) {
|
||||
m, err := New(config.CacheConfig{
|
||||
Enabled: true,
|
||||
TTL: -1 * time.Second, // 立即过期
|
||||
MaxItems: 100,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
m.Set(&models.SearchRequest{Query: "test"}, models.ProviderTavily, &models.SearchResult{})
|
||||
got := m.Get(&models.SearchRequest{Query: "test"}, models.ProviderTavily)
|
||||
if got != nil {
|
||||
t.Fatal("expected nil for expired entry")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCache_Stats(t *testing.T) {
|
||||
m, err := New(config.CacheConfig{
|
||||
Enabled: true,
|
||||
TTL: 5 * time.Minute,
|
||||
MaxItems: 100,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Miss then hit
|
||||
m.Get(&models.SearchRequest{Query: "miss"}, models.ProviderTavily)
|
||||
m.Set(&models.SearchRequest{Query: "hit"}, models.ProviderTavily, &models.SearchResult{})
|
||||
m.Get(&models.SearchRequest{Query: "hit"}, models.ProviderTavily)
|
||||
|
||||
l1Hits, l2Hits, misses := m.Stats()
|
||||
if l1Hits != 1 {
|
||||
t.Errorf("expected 1 L1 hit, got %d", l1Hits)
|
||||
}
|
||||
if l2Hits != 0 {
|
||||
t.Errorf("expected 0 L2 hits, got %d", l2Hits)
|
||||
}
|
||||
if misses != 1 {
|
||||
t.Errorf("expected 1 miss, got %d", misses)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCache_ProviderIsolation(t *testing.T) {
|
||||
m, err := New(config.CacheConfig{
|
||||
Enabled: true,
|
||||
TTL: 5 * time.Minute,
|
||||
MaxItems: 100,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
req := &models.SearchRequest{Query: "test", MaxResults: 5}
|
||||
m.Set(req, models.ProviderTavily, &models.SearchResult{Query: "test"})
|
||||
|
||||
// Same query, different provider should miss
|
||||
got := m.Get(req, models.ProviderBrave)
|
||||
if got != nil {
|
||||
t.Fatal("expected miss for different provider")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
// Package config 管理博闻配置。
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config 顶层配置。
|
||||
type Config struct {
|
||||
Server ServerConfig `yaml:"server"`
|
||||
Providers ProvidersConfig `yaml:"providers"`
|
||||
Balancer BalancerConfig `yaml:"balancer"`
|
||||
Cache CacheConfig `yaml:"cache"`
|
||||
Metrics MetricsConfig `yaml:"metrics"`
|
||||
}
|
||||
|
||||
// ServerConfig 服务端配置。
|
||||
type ServerConfig struct {
|
||||
Port int `yaml:"port"`
|
||||
ReadTimeout time.Duration `yaml:"read_timeout"`
|
||||
WriteTimeout time.Duration `yaml:"write_timeout"`
|
||||
}
|
||||
|
||||
// ProvidersConfig 聚合所有搜索提供商配置。
|
||||
type ProvidersConfig struct {
|
||||
Tavily ProviderConfig `yaml:"tavily"`
|
||||
Brave ProviderConfig `yaml:"brave"`
|
||||
}
|
||||
|
||||
// ProviderConfig 单个搜索提供商配置。
|
||||
type ProviderConfig struct {
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Timeout time.Duration `yaml:"timeout"`
|
||||
MaxResults int `yaml:"max_results"`
|
||||
HealthCheck HealthCheckConf `yaml:"health_check"`
|
||||
// APIKeys 从环境变量注入,不在 YAML 中
|
||||
APIKeys []string `yaml:"-"`
|
||||
}
|
||||
|
||||
// HealthCheckConf 健康检查配置。
|
||||
type HealthCheckConf struct {
|
||||
Interval time.Duration `yaml:"interval"`
|
||||
Timeout time.Duration `yaml:"timeout"`
|
||||
FailureThreshold int `yaml:"failure_threshold"`
|
||||
}
|
||||
|
||||
// BalancerConfig 负载均衡配置。
|
||||
type BalancerConfig struct {
|
||||
Strategy string `yaml:"strategy"`
|
||||
RetryCount int `yaml:"retry_count"`
|
||||
CircuitBreaker CircuitBreakerCfg `yaml:"circuit_breaker"`
|
||||
}
|
||||
|
||||
// CircuitBreakerCfg 熔断器配置。
|
||||
type CircuitBreakerCfg struct {
|
||||
Enabled bool `yaml:"enabled"`
|
||||
FailureThreshold int `yaml:"failure_threshold"`
|
||||
Cooldown time.Duration `yaml:"cooldown"`
|
||||
}
|
||||
|
||||
// CacheConfig 缓存配置。
|
||||
type CacheConfig struct {
|
||||
Enabled bool `yaml:"enabled"`
|
||||
TTL time.Duration `yaml:"ttl"`
|
||||
MaxItems int `yaml:"max_items"`
|
||||
Redis RedisConfig `yaml:"redis"`
|
||||
}
|
||||
|
||||
// RedisConfig Redis 缓存配置(二级缓存层,可开关)。
|
||||
type RedisConfig struct {
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Addr string `yaml:"addr"`
|
||||
Password string `yaml:"password"`
|
||||
DB int `yaml:"db"`
|
||||
}
|
||||
|
||||
// MetricsConfig 监控指标配置。
|
||||
type MetricsConfig struct {
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Port int `yaml:"port"`
|
||||
Path string `yaml:"path"`
|
||||
}
|
||||
|
||||
// envVarKeys 定义各 provider 的 API Key 环境变量名。
|
||||
var envVarKeys = map[string]string{
|
||||
"tavily": "TAVILY_API_KEYS",
|
||||
"brave": "BRAVE_API_KEYS",
|
||||
}
|
||||
|
||||
// Load 加载配置,API Keys 从环境变量读取。
|
||||
// configPath 为可选 YAML 配置文件路径。
|
||||
func Load(configPath string) (*Config, error) {
|
||||
cfg := &Config{}
|
||||
setDefaults(cfg)
|
||||
|
||||
// 从 YAML 文件加载(如存在)
|
||||
if configPath != "" {
|
||||
if err := loadFromFile(configPath, cfg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
// 从环境变量加载 API Keys
|
||||
cfg.Providers.Tavily.APIKeys = parseEnvKeys(envVarKeys["tavily"])
|
||||
cfg.Providers.Brave.APIKeys = parseEnvKeys(envVarKeys["brave"])
|
||||
|
||||
// 从环境变量加载 Redis 配置
|
||||
if env := os.Getenv("MCP_WEB_SEARCH_REDIS_ENABLED"); env != "" {
|
||||
cfg.Cache.Redis.Enabled, _ = strconv.ParseBool(env)
|
||||
}
|
||||
if env := os.Getenv("MCP_WEB_SEARCH_REDIS_ADDR"); env != "" {
|
||||
cfg.Cache.Redis.Addr = env
|
||||
}
|
||||
if env := os.Getenv("MCP_WEB_SEARCH_REDIS_PASSWORD"); env != "" {
|
||||
cfg.Cache.Redis.Password = env
|
||||
}
|
||||
if env := os.Getenv("MCP_WEB_SEARCH_REDIS_DB"); env != "" {
|
||||
if db, err := strconv.Atoi(env); err == nil {
|
||||
cfg.Cache.Redis.DB = db
|
||||
}
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
func setDefaults(cfg *Config) {
|
||||
cfg.Server.Port = 8080
|
||||
cfg.Server.ReadTimeout = 30 * time.Second
|
||||
cfg.Server.WriteTimeout = 30 * time.Second
|
||||
// Provider 默认启用(可通过 YAML 中 enabled: false 关闭)
|
||||
cfg.Providers.Tavily.Enabled = true
|
||||
cfg.Providers.Tavily.Timeout = 10 * time.Second
|
||||
cfg.Providers.Tavily.MaxResults = 10
|
||||
cfg.Providers.Brave.Enabled = true
|
||||
cfg.Providers.Brave.Timeout = 10 * time.Second
|
||||
cfg.Providers.Brave.MaxResults = 10
|
||||
cfg.Balancer.Strategy = "weighted_round_robin"
|
||||
cfg.Balancer.RetryCount = 2
|
||||
cfg.Balancer.CircuitBreaker.Enabled = true
|
||||
cfg.Balancer.CircuitBreaker.FailureThreshold = 5
|
||||
cfg.Balancer.CircuitBreaker.Cooldown = 30 * time.Second
|
||||
cfg.Cache.Enabled = true
|
||||
cfg.Cache.TTL = 5 * time.Minute
|
||||
cfg.Cache.MaxItems = 1000
|
||||
cfg.Cache.Redis.Enabled = false
|
||||
cfg.Cache.Redis.Addr = "127.0.0.1:6379"
|
||||
cfg.Cache.Redis.Password = ""
|
||||
cfg.Cache.Redis.DB = 0
|
||||
cfg.Metrics.Enabled = true
|
||||
cfg.Metrics.Port = 9090
|
||||
cfg.Metrics.Path = "/metrics"
|
||||
}
|
||||
|
||||
func loadFromFile(path string, cfg *Config) error {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read config: %w", err)
|
||||
}
|
||||
return yamlUnmarshal(data, cfg)
|
||||
}
|
||||
|
||||
func parseEnvKeys(envName string) []string {
|
||||
val := os.Getenv(envName)
|
||||
if val == "" {
|
||||
return nil
|
||||
}
|
||||
parts := strings.Split(val, ",")
|
||||
keys := make([]string, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
if trimmed := strings.TrimSpace(p); trimmed != "" {
|
||||
keys = append(keys, trimmed)
|
||||
}
|
||||
}
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/fsnotify/fsnotify"
|
||||
)
|
||||
|
||||
// Reloader 配置文件热重载器。
|
||||
// 监听 YAML 文件变更,通过 onChange 回调通知调用方。
|
||||
type Reloader struct {
|
||||
path string
|
||||
cfg *Config
|
||||
watcher *fsnotify.Watcher
|
||||
onChange func(old, new *Config) error
|
||||
mu sync.Mutex
|
||||
debounce *time.Timer
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// NewReloader 创建并启动配置热重载。
|
||||
// onChange 在配置变更时被调用,接收 old/new 两份配置。
|
||||
func NewReloader(path string, cfg *Config, onChange func(old, new *Config) error) (*Reloader, error) {
|
||||
w, err := fsnotify.NewWatcher()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := w.Add(path); err != nil {
|
||||
w.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
r := &Reloader{
|
||||
path: path,
|
||||
cfg: cfg,
|
||||
watcher: w,
|
||||
onChange: onChange,
|
||||
debounce: time.NewTimer(0),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
|
||||
// 停止初始化的空 timer
|
||||
if !r.debounce.Stop() {
|
||||
<-r.debounce.C
|
||||
}
|
||||
|
||||
go r.loop()
|
||||
slog.Info("配置热重载已启动", "path", path)
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func (r *Reloader) loop() {
|
||||
for {
|
||||
select {
|
||||
case event := <-r.watcher.Events:
|
||||
if event.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Rename) == 0 {
|
||||
continue
|
||||
}
|
||||
slog.Debug("配置文件变更", "event", event.Name)
|
||||
r.debounce.Reset(500 * time.Millisecond)
|
||||
|
||||
case <-r.debounce.C:
|
||||
r.reload()
|
||||
|
||||
case err := <-r.watcher.Errors:
|
||||
slog.Error("配置热重载监听异常", "error", err)
|
||||
|
||||
case <-r.done:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reloader) reload() {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
old := r.cfg
|
||||
// 只重新加载 YAML 部分(API Keys 仍由环境变量管理)
|
||||
updated := *old // 浅拷贝
|
||||
if err := loadFromFile(r.path, &updated); err != nil {
|
||||
slog.Error("配置热重载读取失败", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 保留环境变量注入的值
|
||||
updated.Providers.Tavily.APIKeys = old.Providers.Tavily.APIKeys
|
||||
updated.Providers.Brave.APIKeys = old.Providers.Brave.APIKeys
|
||||
updated.Cache.Redis.Enabled = old.Cache.Redis.Enabled
|
||||
updated.Cache.Redis.Addr = old.Cache.Redis.Addr
|
||||
updated.Cache.Redis.Password = old.Cache.Redis.Password
|
||||
updated.Cache.Redis.DB = old.Cache.Redis.DB
|
||||
|
||||
if err := r.onChange(r.cfg, &updated); err != nil {
|
||||
slog.Error("配置热重载应用失败", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
r.cfg = &updated
|
||||
slog.Info("配置已热重载")
|
||||
}
|
||||
|
||||
// Stop 停止配置热重载。
|
||||
func (r *Reloader) Stop() {
|
||||
close(r.done)
|
||||
r.watcher.Close()
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
package config
|
||||
|
||||
import "gopkg.in/yaml.v3"
|
||||
|
||||
func yamlUnmarshal(data []byte, cfg *Config) error {
|
||||
return yaml.Unmarshal(data, cfg)
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
// Package logger 提供结构化日志配置。
|
||||
package logger
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Config 日志配置。
|
||||
type Config struct {
|
||||
Level slog.Level
|
||||
AddSource bool
|
||||
JSON bool // true = JSONHandler,false = 彩色 TextHandler
|
||||
}
|
||||
|
||||
// New 创建配置好的 Logger。
|
||||
// 开发环境(JSON=false)输出彩色终端文本,生产环境输出 JSON 供采集。
|
||||
func New(cfg Config) *slog.Logger {
|
||||
var h slog.Handler
|
||||
opts := &slog.HandlerOptions{
|
||||
Level: cfg.Level,
|
||||
AddSource: cfg.AddSource,
|
||||
}
|
||||
if cfg.JSON {
|
||||
h = slog.NewJSONHandler(os.Stdout, opts)
|
||||
} else {
|
||||
h = newColorHandler(os.Stdout, opts)
|
||||
}
|
||||
return slog.New(h)
|
||||
}
|
||||
|
||||
// ANSI 颜色代码
|
||||
const (
|
||||
colorReset = "\033[0m"
|
||||
colorRed = "\033[31m"
|
||||
colorGreen = "\033[32m"
|
||||
colorYellow = "\033[33m"
|
||||
colorBlue = "\033[34m"
|
||||
colorPurple = "\033[35m"
|
||||
colorCyan = "\033[36m"
|
||||
colorGray = "\033[37m"
|
||||
)
|
||||
|
||||
// levelColor 日志级别 → ANSI 颜色映射
|
||||
var levelColor = map[slog.Level]string{
|
||||
slog.LevelDebug: colorGray,
|
||||
slog.LevelInfo: colorGreen,
|
||||
slog.LevelWarn: colorYellow,
|
||||
slog.LevelError: colorRed,
|
||||
}
|
||||
|
||||
// levelLabel 日志级别 → 短标签映射
|
||||
var levelLabel = map[slog.Level]string{
|
||||
slog.LevelDebug: "DEBUG",
|
||||
slog.LevelInfo: "INFO",
|
||||
slog.LevelWarn: "WARN",
|
||||
slog.LevelError: "ERROR",
|
||||
}
|
||||
|
||||
// colorHandler 带颜色的 slog.Handler 包装。
|
||||
type colorHandler struct {
|
||||
inner slog.Handler
|
||||
mu sync.Mutex
|
||||
writer io.Writer
|
||||
}
|
||||
|
||||
func newColorHandler(w io.Writer, opts *slog.HandlerOptions) *colorHandler {
|
||||
return &colorHandler{
|
||||
inner: slog.NewTextHandler(w, opts),
|
||||
writer: w,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *colorHandler) Enabled(ctx context.Context, level slog.Level) bool {
|
||||
return h.inner.Enabled(ctx, level)
|
||||
}
|
||||
|
||||
func (h *colorHandler) Handle(ctx context.Context, r slog.Record) error {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
// 时间(使用 -07:00 让 Go 自动填充时区偏移)
|
||||
timeStr := r.Time.Format("2006-01-02T15:04:05.000-07:00")
|
||||
|
||||
// Level 带颜色
|
||||
lvl := r.Level
|
||||
color := levelColor[lvl]
|
||||
label := levelLabel[lvl]
|
||||
|
||||
// 消息
|
||||
msg := r.Message
|
||||
|
||||
// 组装属性
|
||||
var attrs []string
|
||||
r.Attrs(func(a slog.Attr) bool {
|
||||
attrs = append(attrs, fmt.Sprintf("%s=%v", a.Key, a.Value.Any()))
|
||||
return true
|
||||
})
|
||||
|
||||
// 输出格式:[时间] Level颜色(LEVEL) 消息 attr1=val1 attr2=val2
|
||||
var buf strings.Builder
|
||||
buf.WriteString(timeStr)
|
||||
buf.WriteString(" ")
|
||||
buf.WriteString(color)
|
||||
buf.WriteString(fmt.Sprintf("%-5s", label))
|
||||
buf.WriteString(colorReset)
|
||||
buf.WriteString(" ")
|
||||
buf.WriteString(msg)
|
||||
for _, a := range attrs {
|
||||
buf.WriteString(" ")
|
||||
buf.WriteString(a)
|
||||
}
|
||||
buf.WriteString("\n")
|
||||
|
||||
_, err := fmt.Fprint(h.writer, buf.String())
|
||||
return err
|
||||
}
|
||||
|
||||
func (h *colorHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
|
||||
return &colorHandler{
|
||||
inner: h.inner.WithAttrs(attrs),
|
||||
writer: h.writer,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *colorHandler) WithGroup(name string) slog.Handler {
|
||||
return &colorHandler{
|
||||
inner: h.inner.WithGroup(name),
|
||||
writer: h.writer,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// registerAdminRoutes 注册管理端点到 mux。
|
||||
func (h *Handler) registerAdminRoutes(mux *http.ServeMux) {
|
||||
mux.HandleFunc("/admin/provider/", h.handleProviderAdmin)
|
||||
mux.HandleFunc("/admin/cache/clear", h.handleCacheClear)
|
||||
mux.HandleFunc("/admin/keys", h.handleKeys)
|
||||
}
|
||||
|
||||
// handleProviderAdmin 管理 provider 启用/禁用。
|
||||
// PUT /admin/provider/{name}/enable
|
||||
// PUT /admin/provider/{name}/disable
|
||||
func (h *Handler) handleProviderAdmin(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPut {
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
|
||||
// 解析路径:/admin/provider/{name}/{action}
|
||||
parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/admin/provider/"), "/")
|
||||
if len(parts) != 2 {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "路径格式: /admin/provider/{name}/{action}"})
|
||||
return
|
||||
}
|
||||
name := models.ProviderKind(parts[0])
|
||||
action := parts[1]
|
||||
clientIP := requestIP(r)
|
||||
|
||||
if !h.balancer.HasProvider(name) {
|
||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "provider 不存在: " + string(name)})
|
||||
return
|
||||
}
|
||||
|
||||
switch action {
|
||||
case "enable":
|
||||
h.balancer.EnableProvider(name)
|
||||
slog.Info("管理: 启用 provider", "provider", name, "from", clientIP)
|
||||
writeJSON(w, http.StatusOK, map[string]string{"status": "enabled", "provider": string(name)})
|
||||
case "disable":
|
||||
h.balancer.DisableProvider(name)
|
||||
slog.Info("管理: 禁用 provider", "provider", name, "from", clientIP)
|
||||
writeJSON(w, http.StatusOK, map[string]string{"status": "disabled", "provider": string(name)})
|
||||
default:
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "未知操作: " + action + "(可用: enable, disable)"})
|
||||
}
|
||||
}
|
||||
|
||||
// handleCacheClear 清空缓存。
|
||||
// POST /admin/cache/clear
|
||||
func (h *Handler) handleCacheClear(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
h.cache.Clear()
|
||||
slog.Info("管理: 缓存已清空", "from", requestIP(r))
|
||||
writeJSON(w, http.StatusOK, map[string]string{"status": "cleared"})
|
||||
}
|
||||
|
||||
// handleKeys 列出所有密钥状态。
|
||||
// GET /admin/keys
|
||||
func (h *Handler) handleKeys(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
|
||||
type keyInfo struct {
|
||||
Provider string `json:"provider"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
var keys []keyInfo
|
||||
for _, p := range h.providerKinds {
|
||||
keys = append(keys, keyInfo{
|
||||
Provider: string(p),
|
||||
Enabled: h.balancer.IsProviderEnabled(p),
|
||||
})
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"keys": keys,
|
||||
})
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, data interface{}) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
json.NewEncoder(w).Encode(data)
|
||||
}
|
||||
|
||||
// requestIP 从请求中提取客户端 IP。
|
||||
func requestIP(r *http.Request) string {
|
||||
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
|
||||
return fwd
|
||||
}
|
||||
// RemoteAddr 格式 "ip:port",去掉端口
|
||||
if host := r.RemoteAddr; host != "" {
|
||||
for i := len(host) - 1; i >= 0; i-- {
|
||||
if host[i] == ':' {
|
||||
return host[:i]
|
||||
}
|
||||
}
|
||||
return host
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
// Package mcp 实现 MCP 协议的工具注册、调用处理(传输层无关)。
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sort"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/balancer"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/cache"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/metrics"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// ToolDefinition MCP 工具定义。
|
||||
type ToolDefinition struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
InputSchema interface{} `json:"inputSchema"`
|
||||
}
|
||||
|
||||
// CallToolRequest MCP tools/call JSON-RPC 参数。
|
||||
type CallToolRequest struct {
|
||||
Name string `json:"name"`
|
||||
Input json.RawMessage `json:"input,omitempty"`
|
||||
Arguments json.RawMessage `json:"arguments,omitempty"` // MCP 标准字段名
|
||||
}
|
||||
|
||||
// CallToolResult MCP tools/call 响应体。
|
||||
type CallToolResult struct {
|
||||
Content []ContentItem `json:"content"`
|
||||
IsError bool `json:"isError,omitempty"`
|
||||
}
|
||||
|
||||
// ContentItem MCP 内容项。
|
||||
type ContentItem struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
// Handler 搜索工具核心处理逻辑(传输层无关)。
|
||||
type Handler struct {
|
||||
balancer *balancer.LoadBalancer
|
||||
cache *cache.Manager
|
||||
providerKinds []models.ProviderKind
|
||||
}
|
||||
|
||||
func NewHandler(b *balancer.LoadBalancer, c *cache.Manager, providers []models.ProviderKind) *Handler {
|
||||
return &Handler{
|
||||
balancer: b,
|
||||
cache: c,
|
||||
providerKinds: providers,
|
||||
}
|
||||
}
|
||||
|
||||
// ReloadConfig 应用配置热重载变更。
|
||||
// 由 config.Reloader 在文件变更后调用,只处理 YAML 中可动态变更的字段。
|
||||
func (h *Handler) ReloadConfig(oldCfg, newCfg *config.Config) {
|
||||
// Provider 启用/禁用
|
||||
type pc struct {
|
||||
name models.ProviderKind
|
||||
old config.ProviderConfig
|
||||
new config.ProviderConfig
|
||||
}
|
||||
changes := []pc{
|
||||
{models.ProviderTavily, oldCfg.Providers.Tavily, newCfg.Providers.Tavily},
|
||||
{models.ProviderBrave, oldCfg.Providers.Brave, newCfg.Providers.Brave},
|
||||
}
|
||||
for _, c := range changes {
|
||||
if !c.old.Enabled && c.new.Enabled {
|
||||
h.balancer.EnableProvider(c.name)
|
||||
slog.Info("配置热重载: 启用 provider", "name", c.name)
|
||||
} else if c.old.Enabled && !c.new.Enabled {
|
||||
h.balancer.DisableProvider(c.name)
|
||||
slog.Info("配置热重载: 禁用 provider", "name", c.name)
|
||||
}
|
||||
}
|
||||
|
||||
// 缓存 TTL
|
||||
if oldCfg.Cache.TTL != newCfg.Cache.TTL {
|
||||
h.cache.SetTTL(newCfg.Cache.TTL)
|
||||
slog.Info("配置热重载: 缓存 TTL 更新", "old", oldCfg.Cache.TTL, "new", newCfg.Cache.TTL)
|
||||
}
|
||||
|
||||
// 负载均衡重试次数
|
||||
if oldCfg.Balancer.RetryCount != newCfg.Balancer.RetryCount {
|
||||
h.balancer.SetRetryCount(newCfg.Balancer.RetryCount)
|
||||
slog.Info("配置热重载: 重试次数更新", "old", oldCfg.Balancer.RetryCount, "new", newCfg.Balancer.RetryCount)
|
||||
}
|
||||
}
|
||||
|
||||
// ToolDefinition 返回 web_search 工具定义(MCP tools/list 用)。
|
||||
func (h *Handler) ToolDefinition() ToolDefinition {
|
||||
return ToolDefinition{
|
||||
Name: "web_search",
|
||||
Description: "联网搜索,获取最新的网络信息。支持多搜索引擎。",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"query": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "搜索关键词",
|
||||
},
|
||||
"max_results": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"default": 5,
|
||||
"minimum": 1,
|
||||
"maximum": 20,
|
||||
"description": "最大返回结果数",
|
||||
},
|
||||
"mode": map[string]interface{}{
|
||||
"type": "string",
|
||||
"enum": []string{"first", "merge"},
|
||||
"default": "first",
|
||||
"description": "搜索模式:first(首个成功返回,默认)、merge(并发搜索所有引擎,去重合并)",
|
||||
},
|
||||
"providers": map[string]interface{}{
|
||||
"type": "array",
|
||||
"items": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
"description": "优先使用的搜索引擎(默认全部可用引擎)",
|
||||
},
|
||||
},
|
||||
"required": []string{"query"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// HandleSearch 执行搜索(被 stdio、SSE、HTTP 三种传输层共享)。
|
||||
func (h *Handler) HandleSearch(ctx context.Context, req *models.SearchRequest) *CallToolResult {
|
||||
start := time.Now()
|
||||
|
||||
// 校验
|
||||
if req.Query == "" {
|
||||
metrics.SearchErrors.WithLabelValues("empty_query").Inc()
|
||||
return errorResult("搜索关键词不能为空")
|
||||
}
|
||||
if req.MaxResults <= 0 {
|
||||
req.MaxResults = 5
|
||||
}
|
||||
if req.MaxResults > 20 {
|
||||
req.MaxResults = 20
|
||||
}
|
||||
if req.Mode == "" {
|
||||
req.Mode = models.ModeFirst
|
||||
}
|
||||
|
||||
// 查缓存(用空 provider 查跨 provider 缓存)
|
||||
if cached := h.cache.Get(req, ""); cached != nil {
|
||||
metrics.CacheHits.Inc()
|
||||
metrics.SearchDuration.Observe(time.Since(start).Seconds())
|
||||
return successResult(cached)
|
||||
}
|
||||
metrics.CacheMisses.Inc()
|
||||
|
||||
// 执行搜索
|
||||
if req.Mode == models.ModeMerge {
|
||||
return h.mergeSearch(ctx, req, start)
|
||||
}
|
||||
|
||||
result, err := h.balancer.Search(ctx, req)
|
||||
if err != nil {
|
||||
metrics.SearchErrors.WithLabelValues("search_failed").Inc()
|
||||
return errorResult(fmt.Sprintf("搜索失败: %v", err))
|
||||
}
|
||||
|
||||
// 写入缓存
|
||||
h.cache.Set(req, result.Provider, result)
|
||||
metrics.SearchTotal.WithLabelValues(string(result.Provider)).Inc()
|
||||
|
||||
return successResult(result)
|
||||
}
|
||||
|
||||
func successResult(result *models.SearchResult) *CallToolResult {
|
||||
jsonBytes, _ := json.Marshal(result)
|
||||
return &CallToolResult{
|
||||
Content: []ContentItem{
|
||||
{Type: "json", Text: string(jsonBytes)},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func errorResult(msg string) *CallToolResult {
|
||||
errJSON, _ := json.Marshal(map[string]string{"error": msg})
|
||||
return &CallToolResult{
|
||||
IsError: true,
|
||||
Content: []ContentItem{
|
||||
{Type: "text", Text: string(errJSON)},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// mergeSearch 并发搜索所有 provider,合并去重后返回。
|
||||
func (h *Handler) mergeSearch(ctx context.Context, req *models.SearchRequest, start time.Time) *CallToolResult {
|
||||
results := h.balancer.SearchAll(ctx, req)
|
||||
if len(results) == 0 {
|
||||
metrics.SearchErrors.WithLabelValues("merge_no_results").Inc()
|
||||
return errorResult("所有搜索引擎均未返回结果")
|
||||
}
|
||||
|
||||
// 合并 + URL 去重
|
||||
seen := make(map[string]bool)
|
||||
var merged []models.ResultItem
|
||||
for _, res := range results {
|
||||
for _, item := range res.Results {
|
||||
if seen[item.URL] {
|
||||
continue
|
||||
}
|
||||
seen[item.URL] = true
|
||||
merged = append(merged, item)
|
||||
}
|
||||
}
|
||||
|
||||
// 按 score 降序排序(无 score 的排最后)
|
||||
sort.Slice(merged, func(i, j int) bool {
|
||||
return merged[i].Score > merged[j].Score
|
||||
})
|
||||
|
||||
// 截断到 max_results
|
||||
if len(merged) > req.MaxResults {
|
||||
merged = merged[:req.MaxResults]
|
||||
}
|
||||
|
||||
final := &models.SearchResult{
|
||||
Provider: models.ProviderKind("merge"),
|
||||
Query: req.Query,
|
||||
Results: merged,
|
||||
TotalCount: len(merged),
|
||||
DurationMs: time.Since(start).Milliseconds(),
|
||||
Cached: false,
|
||||
}
|
||||
|
||||
h.cache.Set(req, "", final)
|
||||
return successResult(final)
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/cache"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
func TestHandler_ToolDefinition(t *testing.T) {
|
||||
h := NewHandler(nil, nil, nil)
|
||||
def := h.ToolDefinition()
|
||||
if def.Name != "web_search" {
|
||||
t.Errorf("expected 'web_search', got '%s'", def.Name)
|
||||
}
|
||||
if def.Description == "" {
|
||||
t.Error("expected non-empty description")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_HandleSearch_EmptyQuery(t *testing.T) {
|
||||
c, _ := cache.New(config.CacheConfig{Enabled: false})
|
||||
h := NewHandler(nil, c, nil)
|
||||
result := h.HandleSearch(context.Background(), &models.SearchRequest{Query: ""})
|
||||
if !result.IsError {
|
||||
t.Fatal("expected error for empty query")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandler_HandleSearch_DefaultMaxResults(t *testing.T) {
|
||||
req := &models.SearchRequest{Query: "test"}
|
||||
if req.MaxResults <= 0 {
|
||||
req.MaxResults = 5
|
||||
}
|
||||
if req.MaxResults > 20 {
|
||||
req.MaxResults = 20
|
||||
}
|
||||
if req.MaxResults != 5 {
|
||||
t.Errorf("expected default 5, got %d", req.MaxResults)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// HandleMonitor 返回 HTML 监控面板(GET /)。
|
||||
func (h *Handler) HandleMonitor(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
|
||||
uptime := time.Since(startTime).Round(time.Second).String()
|
||||
|
||||
// Provider 表格行
|
||||
var providerRows string
|
||||
for _, p := range h.providerKinds {
|
||||
counts := h.balancer.ActiveKeyCount()
|
||||
activeKeys := counts[p]
|
||||
enabled := h.balancer.IsProviderEnabled(p)
|
||||
statusIcon := "🟢"
|
||||
statusText := "启用"
|
||||
if !enabled {
|
||||
statusIcon = "🔴"
|
||||
statusText = "禁用"
|
||||
}
|
||||
providerRows += fmt.Sprintf(`<tr>
|
||||
<td>%s</td>
|
||||
<td><span class="badge badge-provider">%s</span></td>
|
||||
<td>%s <span class="status-text">%s</span></td>
|
||||
<td class="num">%d</td>
|
||||
</tr>`, statusIcon, p, statusIcon, statusText, activeKeys)
|
||||
}
|
||||
|
||||
// 缓存统计
|
||||
l1Hits, l2Hits, misses := h.cache.Stats()
|
||||
hitRate := h.cache.HitRate()
|
||||
|
||||
html := fmt.Sprintf(`<!DOCTYPE html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>博闻 · 监控面板</title>
|
||||
<style>
|
||||
* { margin: 0; padding: 0; box-sizing: border-box; }
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Noto Sans SC", sans-serif;
|
||||
background: #0f172a;
|
||||
color: #e2e8f0;
|
||||
padding: 24px;
|
||||
line-height: 1.6;
|
||||
}
|
||||
.container { max-width: 960px; margin: 0 auto; }
|
||||
h1 {
|
||||
font-size: 24px;
|
||||
font-weight: 600;
|
||||
margin-bottom: 4px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 12px;
|
||||
}
|
||||
h1 small {
|
||||
font-size: 14px;
|
||||
font-weight: 400;
|
||||
color: #64748b;
|
||||
}
|
||||
.subtitle {
|
||||
color: #64748b;
|
||||
font-size: 14px;
|
||||
margin-bottom: 24px;
|
||||
}
|
||||
.grid {
|
||||
display: grid;
|
||||
grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
|
||||
gap: 16px;
|
||||
margin-bottom: 24px;
|
||||
}
|
||||
.card {
|
||||
background: #1e293b;
|
||||
border-radius: 12px;
|
||||
padding: 16px 20px;
|
||||
border: 1px solid #334155;
|
||||
}
|
||||
.card .label {
|
||||
font-size: 12px;
|
||||
color: #64748b;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.05em;
|
||||
margin-bottom: 4px;
|
||||
}
|
||||
.card .value {
|
||||
font-size: 28px;
|
||||
font-weight: 700;
|
||||
color: #38bdf8;
|
||||
}
|
||||
.card .value.green { color: #4ade80; }
|
||||
.card .value.yellow { color: #fbbf24; }
|
||||
.card .value.red { color: #f87171; }
|
||||
.section-title {
|
||||
font-size: 16px;
|
||||
font-weight: 600;
|
||||
margin-bottom: 12px;
|
||||
color: #94a3b8;
|
||||
}
|
||||
table {
|
||||
width: 100%%;
|
||||
border-collapse: collapse;
|
||||
background: #1e293b;
|
||||
border-radius: 12px;
|
||||
overflow: hidden;
|
||||
border: 1px solid #334155;
|
||||
margin-bottom: 24px;
|
||||
}
|
||||
th {
|
||||
text-align: left;
|
||||
padding: 12px 16px;
|
||||
font-size: 12px;
|
||||
color: #64748b;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.05em;
|
||||
background: #1a2332;
|
||||
border-bottom: 1px solid #334155;
|
||||
}
|
||||
td {
|
||||
padding: 10px 16px;
|
||||
border-bottom: 1px solid #1e293b;
|
||||
font-size: 14px;
|
||||
}
|
||||
tr:last-child td { border-bottom: none; }
|
||||
tr:hover { background: #243348; }
|
||||
.num { text-align: right; font-variant-numeric: tabular-nums; }
|
||||
.badge {
|
||||
display: inline-block;
|
||||
padding: 2px 8px;
|
||||
border-radius: 4px;
|
||||
font-size: 12px;
|
||||
font-weight: 500;
|
||||
}
|
||||
.badge-provider {
|
||||
background: #1e3a5f;
|
||||
color: #93c5fd;
|
||||
}
|
||||
.status-text {
|
||||
font-size: 13px;
|
||||
color: #94a3b8;
|
||||
margin-left: 4px;
|
||||
}
|
||||
a { color: #38bdf8; text-decoration: none; }
|
||||
a:hover { text-decoration: underline; }
|
||||
.links { margin-top: 16px; display: flex; gap: 16px; font-size: 14px; }
|
||||
.footer { margin-top: 32px; padding-top: 16px; border-top: 1px solid #334155; color: #475569; font-size: 12px; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<h1>🔍 博闻 <small>Bowen Web Search</small></h1>
|
||||
<div class="subtitle">运行时间 %s</div>
|
||||
|
||||
<div class="grid">
|
||||
<div class="card">
|
||||
<div class="label">Provider</div>
|
||||
<div class="value">%d</div>
|
||||
</div>
|
||||
<div class="card">
|
||||
<div class="label">活跃密钥</div>
|
||||
<div class="value">%d</div>
|
||||
</div>
|
||||
<div class="card">
|
||||
<div class="label">缓存命中率</div>
|
||||
<div class="value %s">%d%%%%</div>
|
||||
</div>
|
||||
<div class="card">
|
||||
<div class="label">缓存条目</div>
|
||||
<div class="value">%d</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="section-title">📡 Provider 状态</div>
|
||||
<table>
|
||||
<thead>
|
||||
<tr><th>状态</th><th>名称</th><th>开关</th><th class="num">活跃密钥</th></tr>
|
||||
</thead>
|
||||
<tbody>%s</tbody>
|
||||
</table>
|
||||
|
||||
<div class="section-title">💾 缓存统计</div>
|
||||
<table>
|
||||
<thead>
|
||||
<tr><th>指标</th><th class="num">值</th></tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr><td>L1 命中(内存 LRU)</td><td class="num">%d</td></tr>
|
||||
<tr><td>L2 命中(Redis)</td><td class="num">%d</td></tr>
|
||||
<tr><td>未命中</td><td class="num">%d</td></tr>
|
||||
<tr><td>总命中率</td><td class="num">%.1f%%%%</td></tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
<div class="links">
|
||||
<span>🔗</span>
|
||||
<a href="/status">/status (JSON)</a>
|
||||
<a href="/admin/keys">/admin/keys</a>
|
||||
<a href="/health">/health</a>
|
||||
</div>
|
||||
|
||||
<div class="footer">博闻 v2.0.0 · 数据实时 · 页面每 30 秒刷新</div>
|
||||
</div>
|
||||
<script>setTimeout(function(){ location.reload(); }, 30000);</script>
|
||||
</body>
|
||||
</html>`,
|
||||
uptime,
|
||||
len(h.providerKinds),
|
||||
sumActiveKeys(h.balancer.ActiveKeyCount()),
|
||||
hitRateColor(hitRate),
|
||||
int(hitRate*100),
|
||||
l1Hits+l2Hits+misses,
|
||||
providerRows,
|
||||
l1Hits, l2Hits, misses, hitRate*100,
|
||||
)
|
||||
|
||||
fmt.Fprint(w, html)
|
||||
}
|
||||
|
||||
func sumActiveKeys(m map[models.ProviderKind]int) int {
|
||||
s := 0
|
||||
for _, v := range m {
|
||||
s += v
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func hitRateColor(rate float64) string {
|
||||
if rate >= 0.6 {
|
||||
return "green"
|
||||
} else if rate >= 0.3 {
|
||||
return "yellow"
|
||||
}
|
||||
return "red"
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
// LoggingMiddleware 记录 HTTP 请求的源 IP、方法、路径和处理耗时。
|
||||
func LoggingMiddleware(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
start := time.Now()
|
||||
ip := r.RemoteAddr
|
||||
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
|
||||
ip = fwd
|
||||
}
|
||||
method := r.Method
|
||||
path := r.URL.Path
|
||||
next.ServeHTTP(w, r)
|
||||
slog.Info("HTTP",
|
||||
"method", method,
|
||||
"path", path,
|
||||
"from", ip,
|
||||
"ms", time.Since(start).Milliseconds(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
// ServeHTTP 启动 Streamable HTTP 传输模式的 MCP 服务。
|
||||
//
|
||||
// 传输协议:MCP Streamable HTTP(JSON-RPC over HTTP POST /)
|
||||
// POST / → MCP 处理,GET / → 帮助页
|
||||
func (h *Handler) ServeHTTP(addr string) error {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/", h.handleRoot)
|
||||
// 管理端点
|
||||
mux.HandleFunc("/health", h.HandleHealth)
|
||||
mux.HandleFunc("/status", h.handleStatus)
|
||||
h.registerAdminRoutes(mux)
|
||||
|
||||
slog.Info("博闻 HTTP 服务启动", "addr", addr, "providers", h.providerKinds)
|
||||
return http.ListenAndServe(addr, LoggingMiddleware(mux))
|
||||
}
|
||||
|
||||
// handleRoot MCP 服务根路径。
|
||||
// POST → MCP Streamable HTTP(JSON-RPC 2.0)
|
||||
// GET → 帮助页
|
||||
func (h *Handler) handleRoot(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodPost {
|
||||
h.handleStreamableHTTP(w, r)
|
||||
return
|
||||
}
|
||||
if r.URL.Path != "/" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
||||
fmt.Fprint(w, `博闻 Bowen · MCP Web Search
|
||||
|
||||
传输协议:
|
||||
POST / MCP Streamable HTTP(JSON-RPC 2.0,单端点)
|
||||
|
||||
管理端点:
|
||||
GET /health 存活探针
|
||||
GET /status 系统状态
|
||||
GET /admin/keys 密钥状态
|
||||
PUT /admin/provider/:name/enable 启用 provider
|
||||
PUT /admin/provider/:name/disable 禁用 provider
|
||||
POST /admin/cache/clear 清空缓存
|
||||
|
||||
监控面板(metrics 端口):
|
||||
GET / HTML 仪表盘
|
||||
GET /metrics Prometheus 原始指标
|
||||
|
||||
Stdio 模式:go run ./cmd/server --stdio
|
||||
`)
|
||||
}
|
||||
|
||||
// handleStreamableHTTP MCP Streamable HTTP 传输端点。
|
||||
// 接收 JSON-RPC 2.0 POST 请求,返回 JSON-RPC 响应。
|
||||
// 现阶段返回完整响应(非流式),后续可扩展为 text/event-stream 流式传输。
|
||||
func (h *Handler) handleStreamableHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
|
||||
var req JSONRPCRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
json.NewEncoder(w).Encode(jsonRPCError(nil, -32700, "Parse error"))
|
||||
return
|
||||
}
|
||||
|
||||
resp := h.handleJSONRPC(req)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SystemStatus 博闻系统状态(通过 /status 或管理 API 查询)。
|
||||
type SystemStatus struct {
|
||||
Version string `json:"version"`
|
||||
Uptime string `json:"uptime"`
|
||||
Providers []ProviderStatus `json:"providers"`
|
||||
Cache CacheStats `json:"cache"`
|
||||
}
|
||||
|
||||
// ProviderStatus provider 状态摘要。
|
||||
type ProviderStatus struct {
|
||||
Name string `json:"name"`
|
||||
ActiveKeys int `json:"active_keys"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
// CacheStats 缓存状态摘要。
|
||||
type CacheStats struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
L1Hits int64 `json:"l1_hits"`
|
||||
L2Hits int64 `json:"l2_hits"`
|
||||
Misses int64 `json:"misses"`
|
||||
HitRate float64 `json:"hit_rate"`
|
||||
}
|
||||
|
||||
// startTime 记录博闻启动时间,用于 /status 的 uptime 计算。
|
||||
var startTime = time.Now()
|
||||
|
||||
// HandleHealth 存活探针(GET /health → 200)。
|
||||
func (h *Handler) HandleHealth(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
|
||||
}
|
||||
|
||||
// handleStatus 系统状态(GET /status → SystemStatus)。
|
||||
func (h *Handler) handleStatus(w http.ResponseWriter, r *http.Request) {
|
||||
// Provider 状态
|
||||
var providers []ProviderStatus
|
||||
for _, p := range h.providerKinds {
|
||||
counts := h.balancer.ActiveKeyCount()
|
||||
activeKeys := counts[p]
|
||||
providers = append(providers, ProviderStatus{
|
||||
Name: string(p),
|
||||
ActiveKeys: activeKeys,
|
||||
Enabled: activeKeys > 0,
|
||||
})
|
||||
}
|
||||
|
||||
// 缓存状态
|
||||
l1Hits, l2Hits, misses := h.cache.Stats()
|
||||
cacheStats := CacheStats{
|
||||
Enabled: true,
|
||||
L1Hits: l1Hits,
|
||||
L2Hits: l2Hits,
|
||||
Misses: misses,
|
||||
HitRate: h.cache.HitRate(),
|
||||
}
|
||||
|
||||
status := SystemStatus{
|
||||
Version: "2.0.0",
|
||||
Uptime: time.Since(startTime).Round(time.Second).String(),
|
||||
Providers: providers,
|
||||
Cache: cacheStats,
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(status)
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
|
||||
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
||||
)
|
||||
|
||||
// JSONRPCRequest JSON-RPC 2.0 请求。
|
||||
type JSONRPCRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID *int `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params json.RawMessage `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
// JSONRPCResponse JSON-RPC 2.0 响应。
|
||||
type JSONRPCResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID *int `json:"id"`
|
||||
Result interface{} `json:"result,omitempty"`
|
||||
Error *RPCError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// RPCError JSON-RPC 错误。
|
||||
type RPCError struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// ServeStdio 通过 stdin/stdout 提供 MCP JSON-RPC 服务。
|
||||
func (h *Handler) ServeStdio() {
|
||||
slog.Info("博闻 stdio 模式启动", "providers", h.providerKinds)
|
||||
reader := bufio.NewReader(os.Stdin)
|
||||
encoder := json.NewEncoder(os.Stdout)
|
||||
|
||||
for {
|
||||
line, err := reader.ReadString('\n')
|
||||
if err != nil {
|
||||
slog.Error("stdio read error", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
var req JSONRPCRequest
|
||||
if err := json.Unmarshal([]byte(line), &req); err != nil {
|
||||
slog.Warn("stdio parse error", "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
resp := h.handleJSONRPC(req)
|
||||
if err := encoder.Encode(resp); err != nil {
|
||||
slog.Error("stdio write error", "error", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) handleJSONRPC(req JSONRPCRequest) JSONRPCResponse {
|
||||
switch req.Method {
|
||||
case "tools/list":
|
||||
return JSONRPCResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: req.ID,
|
||||
Result: map[string]interface{}{
|
||||
"tools": []ToolDefinition{h.ToolDefinition()},
|
||||
},
|
||||
}
|
||||
|
||||
case "tools/call":
|
||||
var callReq CallToolRequest
|
||||
if err := json.Unmarshal(req.Params, &callReq); err != nil {
|
||||
return jsonRPCError(req.ID, -32602, "Invalid params")
|
||||
}
|
||||
if callReq.Name != "web_search" {
|
||||
return jsonRPCError(req.ID, -32601, fmt.Sprintf("Tool not found: %s", callReq.Name))
|
||||
}
|
||||
|
||||
// 优先使用 MCP 标准 arguments 字段,兼容旧 input 字段
|
||||
rawInput := callReq.Arguments
|
||||
if len(rawInput) == 0 {
|
||||
rawInput = callReq.Input
|
||||
}
|
||||
|
||||
var searchReq models.SearchRequest
|
||||
if err := json.Unmarshal(rawInput, &searchReq); err != nil {
|
||||
return jsonRPCError(req.ID, -32602, "Invalid search request parameters: "+err.Error())
|
||||
}
|
||||
|
||||
result := h.HandleSearch(context.Background(), &searchReq)
|
||||
return JSONRPCResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: req.ID,
|
||||
Result: result,
|
||||
}
|
||||
|
||||
default:
|
||||
return jsonRPCError(req.ID, -32601, fmt.Sprintf("Method not found: %s", req.Method))
|
||||
}
|
||||
}
|
||||
|
||||
func jsonRPCError(id *int, code int, msg string) JSONRPCResponse {
|
||||
return JSONRPCResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Error: &RPCError{Code: code, Message: msg},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
// Package metrics 定义 Prometheus 监控指标。
|
||||
package metrics
|
||||
|
||||
import (
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/prometheus/client_golang/prometheus/promauto"
|
||||
)
|
||||
|
||||
var (
|
||||
// SearchTotal 按 provider 统计搜索请求总数。
|
||||
SearchTotal = promauto.NewCounterVec(
|
||||
prometheus.CounterOpts{
|
||||
Name: "bowen_search_total",
|
||||
Help: "Total search requests by provider.",
|
||||
},
|
||||
[]string{"provider"},
|
||||
)
|
||||
|
||||
// SearchDuration 搜索请求耗时直方图。
|
||||
SearchDuration = promauto.NewHistogram(
|
||||
prometheus.HistogramOpts{
|
||||
Name: "bowen_search_duration_seconds",
|
||||
Help: "Search request duration in seconds.",
|
||||
Buckets: prometheus.DefBuckets,
|
||||
},
|
||||
)
|
||||
|
||||
// SearchErrors 按原因统计搜索错误。
|
||||
SearchErrors = promauto.NewCounterVec(
|
||||
prometheus.CounterOpts{
|
||||
Name: "bowen_search_errors_total",
|
||||
Help: "Search errors by reason.",
|
||||
},
|
||||
[]string{"reason"},
|
||||
)
|
||||
|
||||
// CacheHits 缓存命中总数。
|
||||
CacheHits = promauto.NewCounter(
|
||||
prometheus.CounterOpts{
|
||||
Name: "bowen_cache_hits_total",
|
||||
Help: "Total cache hits.",
|
||||
},
|
||||
)
|
||||
|
||||
// CacheMisses 缓存未命中总数。
|
||||
CacheMisses = promauto.NewCounter(
|
||||
prometheus.CounterOpts{
|
||||
Name: "bowen_cache_misses_total",
|
||||
Help: "Total cache misses.",
|
||||
},
|
||||
)
|
||||
|
||||
// CacheHitsL2 Redis 二级缓存命中数。
|
||||
CacheHitsL2 = promauto.NewCounter(
|
||||
prometheus.CounterOpts{
|
||||
Name: "bowen_cache_hits_l2_total",
|
||||
Help: "L2 (Redis) cache hits.",
|
||||
},
|
||||
)
|
||||
|
||||
// ActiveKeys 各 provider 的活跃密钥数。
|
||||
ActiveKeys = promauto.NewGaugeVec(
|
||||
prometheus.GaugeOpts{
|
||||
Name: "bowen_active_keys",
|
||||
Help: "Number of active API keys by provider.",
|
||||
},
|
||||
[]string{"provider"},
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,45 @@
|
||||
// Package models 定义博闻的数据类型。
|
||||
package models
|
||||
|
||||
// ProviderKind 搜索提供商枚举。
|
||||
type ProviderKind string
|
||||
|
||||
const (
|
||||
ProviderTavily ProviderKind = "tavily"
|
||||
ProviderBrave ProviderKind = "brave"
|
||||
)
|
||||
|
||||
// SearchMode 搜索结果聚合模式。
|
||||
type SearchMode string
|
||||
|
||||
const (
|
||||
ModeFirst SearchMode = "first" // 首个成功返回(低延迟,默认)
|
||||
ModeMerge SearchMode = "merge" // 合并多 provider 结果(高召回)
|
||||
)
|
||||
|
||||
// SearchRequest MCP 工具入参。
|
||||
type SearchRequest struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
Providers []string `json:"providers,omitempty"`
|
||||
Mode SearchMode `json:"mode,omitempty"`
|
||||
}
|
||||
|
||||
// SearchResult 统一搜索响应。
|
||||
type SearchResult struct {
|
||||
Provider ProviderKind `json:"provider"`
|
||||
Query string `json:"query"`
|
||||
Results []ResultItem `json:"results"`
|
||||
TotalCount int `json:"total_count"`
|
||||
DurationMs int64 `json:"duration_ms"`
|
||||
Cached bool `json:"cached"`
|
||||
}
|
||||
|
||||
// ResultItem 单条搜索结果。
|
||||
type ResultItem struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Snippet string `json:"snippet"`
|
||||
PublishedAt string `json:"published_at,omitempty"`
|
||||
Score float64 `json:"score,omitempty"`
|
||||
}
|
||||
@@ -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