support proxy in web tool
This commit is contained in:
parent
100356e8ec
commit
3ab5b306b7
4 changed files with 90 additions and 8 deletions
|
|
@ -106,6 +106,7 @@ func registerSharedTools(
|
||||||
PerplexityAPIKey: cfg.Tools.Web.Perplexity.APIKey,
|
PerplexityAPIKey: cfg.Tools.Web.Perplexity.APIKey,
|
||||||
PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults,
|
PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults,
|
||||||
PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled,
|
PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled,
|
||||||
|
Proxy: cfg.Tools.Web.Proxy,
|
||||||
}); searchTool != nil {
|
}); searchTool != nil {
|
||||||
agent.Tools.Register(searchTool)
|
agent.Tools.Register(searchTool)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -453,6 +453,7 @@ type WebToolsConfig struct {
|
||||||
Tavily TavilyConfig `json:"tavily"`
|
Tavily TavilyConfig `json:"tavily"`
|
||||||
DuckDuckGo DuckDuckGoConfig `json:"duckduckgo"`
|
DuckDuckGo DuckDuckGoConfig `json:"duckduckgo"`
|
||||||
Perplexity PerplexityConfig `json:"perplexity"`
|
Perplexity PerplexityConfig `json:"perplexity"`
|
||||||
|
Proxy string `json:"proxy,omitempty" env:"PICOCLAW_TOOLS_WEB_PROXY"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type CronToolsConfig struct {
|
type CronToolsConfig struct {
|
||||||
|
|
|
||||||
|
|
@ -277,6 +277,7 @@ func DefaultConfig() *Config {
|
||||||
},
|
},
|
||||||
Tools: ToolsConfig{
|
Tools: ToolsConfig{
|
||||||
Web: WebToolsConfig{
|
Web: WebToolsConfig{
|
||||||
|
Proxy: "",
|
||||||
Brave: BraveConfig{
|
Brave: BraveConfig{
|
||||||
Enabled: false,
|
Enabled: false,
|
||||||
APIKey: "",
|
APIKey: "",
|
||||||
|
|
|
||||||
|
|
@ -6,11 +6,14 @@ import (
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/url"
|
"net/url"
|
||||||
"regexp"
|
"regexp"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"golang.org/x/net/proxy"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
|
|
@ -23,6 +26,7 @@ type SearchProvider interface {
|
||||||
|
|
||||||
type BraveSearchProvider struct {
|
type BraveSearchProvider struct {
|
||||||
apiKey string
|
apiKey string
|
||||||
|
proxy string
|
||||||
}
|
}
|
||||||
|
|
||||||
func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
||||||
|
|
@ -37,7 +41,10 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
|
||||||
req.Header.Set("Accept", "application/json")
|
req.Header.Set("Accept", "application/json")
|
||||||
req.Header.Set("X-Subscription-Token", p.apiKey)
|
req.Header.Set("X-Subscription-Token", p.apiKey)
|
||||||
|
|
||||||
client := &http.Client{Timeout: 10 * time.Second}
|
client, err := clientForProxy(p.proxy, 10*time.Second)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to create http client: %w", err)
|
||||||
|
}
|
||||||
resp, err := client.Do(req)
|
resp, err := client.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("request failed: %w", err)
|
return "", fmt.Errorf("request failed: %w", err)
|
||||||
|
|
@ -88,6 +95,7 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
|
||||||
type TavilySearchProvider struct {
|
type TavilySearchProvider struct {
|
||||||
apiKey string
|
apiKey string
|
||||||
baseURL string
|
baseURL string
|
||||||
|
proxy string
|
||||||
}
|
}
|
||||||
|
|
||||||
func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
||||||
|
|
@ -119,7 +127,10 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
|
||||||
req.Header.Set("Content-Type", "application/json")
|
req.Header.Set("Content-Type", "application/json")
|
||||||
req.Header.Set("User-Agent", userAgent)
|
req.Header.Set("User-Agent", userAgent)
|
||||||
|
|
||||||
client := &http.Client{Timeout: 10 * time.Second}
|
client, err := clientForProxy(p.proxy, 10*time.Second)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to create http client: %w", err)
|
||||||
|
}
|
||||||
resp, err := client.Do(req)
|
resp, err := client.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("request failed: %w", err)
|
return "", fmt.Errorf("request failed: %w", err)
|
||||||
|
|
@ -167,7 +178,10 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
|
||||||
return strings.Join(lines, "\n"), nil
|
return strings.Join(lines, "\n"), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
type DuckDuckGoSearchProvider struct{}
|
// duckduckgo provider may use an optional proxy (socks5 or http)
|
||||||
|
type DuckDuckGoSearchProvider struct {
|
||||||
|
proxy string
|
||||||
|
}
|
||||||
|
|
||||||
func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
||||||
searchURL := fmt.Sprintf("https://html.duckduckgo.com/html/?q=%s", url.QueryEscape(query))
|
searchURL := fmt.Sprintf("https://html.duckduckgo.com/html/?q=%s", url.QueryEscape(query))
|
||||||
|
|
@ -179,7 +193,10 @@ func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, cou
|
||||||
|
|
||||||
req.Header.Set("User-Agent", userAgent)
|
req.Header.Set("User-Agent", userAgent)
|
||||||
|
|
||||||
client := &http.Client{Timeout: 10 * time.Second}
|
client, err := clientForProxy(p.proxy, 10*time.Second)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to create http client: %w", err)
|
||||||
|
}
|
||||||
resp, err := client.Do(req)
|
resp, err := client.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("request failed: %w", err)
|
return "", fmt.Errorf("request failed: %w", err)
|
||||||
|
|
@ -259,8 +276,64 @@ func stripTags(content string) string {
|
||||||
return re.ReplaceAllString(content, "")
|
return re.ReplaceAllString(content, "")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// clientForProxy returns an *http.Client honoring the provided proxy string.
|
||||||
|
// Supported formats:
|
||||||
|
// - "socks5://host:port" or plain "host:port" (treated as socks5)
|
||||||
|
// - any URL with scheme (http://, https://) will be used as an HTTP proxy via ProxyURL
|
||||||
|
func clientForProxy(proxyStr string, timeout time.Duration) (*http.Client, error) {
|
||||||
|
transport := &http.Transport{
|
||||||
|
TLSHandshakeTimeout: 15 * time.Second,
|
||||||
|
MaxIdleConns: 10,
|
||||||
|
IdleConnTimeout: 30 * time.Second,
|
||||||
|
DisableCompression: false,
|
||||||
|
}
|
||||||
|
|
||||||
|
if proxyStr == "" {
|
||||||
|
return &http.Client{Timeout: timeout, Transport: transport}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// If no scheme provided, assume socks5
|
||||||
|
if !strings.Contains(proxyStr, "://") {
|
||||||
|
// treat as socks5 host:port
|
||||||
|
dialer, err := proxy.SOCKS5("tcp", proxyStr, nil, proxy.Direct)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||||
|
return dialer.Dial(network, addr)
|
||||||
|
}
|
||||||
|
return &http.Client{Timeout: timeout, Transport: transport}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
u, err := url.Parse(proxyStr)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
switch u.Scheme {
|
||||||
|
case "socks5", "socks5h":
|
||||||
|
var auth *proxy.Auth
|
||||||
|
if u.User != nil {
|
||||||
|
password, _ := u.User.Password()
|
||||||
|
auth = &proxy.Auth{User: u.User.Username(), Password: password}
|
||||||
|
}
|
||||||
|
dialer, err := proxy.SOCKS5("tcp", u.Host, auth, proxy.Direct)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||||
|
return dialer.Dial(network, addr)
|
||||||
|
}
|
||||||
|
return &http.Client{Timeout: timeout, Transport: transport}, nil
|
||||||
|
default:
|
||||||
|
transport.Proxy = http.ProxyURL(u)
|
||||||
|
return &http.Client{Timeout: timeout, Transport: transport}, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
type PerplexitySearchProvider struct {
|
type PerplexitySearchProvider struct {
|
||||||
apiKey string
|
apiKey string
|
||||||
|
proxy string
|
||||||
}
|
}
|
||||||
|
|
||||||
func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
||||||
|
|
@ -295,7 +368,11 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou
|
||||||
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
||||||
req.Header.Set("User-Agent", userAgent)
|
req.Header.Set("User-Agent", userAgent)
|
||||||
|
|
||||||
client := &http.Client{Timeout: 30 * time.Second}
|
client, err := clientForProxy(p.proxy, 30*time.Second)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to create http client: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
resp, err := client.Do(req)
|
resp, err := client.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("request failed: %w", err)
|
return "", fmt.Errorf("request failed: %w", err)
|
||||||
|
|
@ -348,6 +425,7 @@ type WebSearchToolOptions struct {
|
||||||
PerplexityAPIKey string
|
PerplexityAPIKey string
|
||||||
PerplexityMaxResults int
|
PerplexityMaxResults int
|
||||||
PerplexityEnabled bool
|
PerplexityEnabled bool
|
||||||
|
Proxy string
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewWebSearchTool(opts WebSearchToolOptions) *WebSearchTool {
|
func NewWebSearchTool(opts WebSearchToolOptions) *WebSearchTool {
|
||||||
|
|
@ -356,12 +434,12 @@ func NewWebSearchTool(opts WebSearchToolOptions) *WebSearchTool {
|
||||||
|
|
||||||
// Priority: Perplexity > Brave > Tavily > DuckDuckGo
|
// Priority: Perplexity > Brave > Tavily > DuckDuckGo
|
||||||
if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" {
|
if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" {
|
||||||
provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey}
|
provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy}
|
||||||
if opts.PerplexityMaxResults > 0 {
|
if opts.PerplexityMaxResults > 0 {
|
||||||
maxResults = opts.PerplexityMaxResults
|
maxResults = opts.PerplexityMaxResults
|
||||||
}
|
}
|
||||||
} else if opts.BraveEnabled && opts.BraveAPIKey != "" {
|
} else if opts.BraveEnabled && opts.BraveAPIKey != "" {
|
||||||
provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey}
|
provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy}
|
||||||
if opts.BraveMaxResults > 0 {
|
if opts.BraveMaxResults > 0 {
|
||||||
maxResults = opts.BraveMaxResults
|
maxResults = opts.BraveMaxResults
|
||||||
}
|
}
|
||||||
|
|
@ -369,12 +447,13 @@ func NewWebSearchTool(opts WebSearchToolOptions) *WebSearchTool {
|
||||||
provider = &TavilySearchProvider{
|
provider = &TavilySearchProvider{
|
||||||
apiKey: opts.TavilyAPIKey,
|
apiKey: opts.TavilyAPIKey,
|
||||||
baseURL: opts.TavilyBaseURL,
|
baseURL: opts.TavilyBaseURL,
|
||||||
|
proxy: opts.Proxy,
|
||||||
}
|
}
|
||||||
if opts.TavilyMaxResults > 0 {
|
if opts.TavilyMaxResults > 0 {
|
||||||
maxResults = opts.TavilyMaxResults
|
maxResults = opts.TavilyMaxResults
|
||||||
}
|
}
|
||||||
} else if opts.DuckDuckGoEnabled {
|
} else if opts.DuckDuckGoEnabled {
|
||||||
provider = &DuckDuckGoSearchProvider{}
|
provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy}
|
||||||
if opts.DuckDuckGoMaxResults > 0 {
|
if opts.DuckDuckGoMaxResults > 0 {
|
||||||
maxResults = opts.DuckDuckGoMaxResults
|
maxResults = opts.DuckDuckGoMaxResults
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue