WeiCHEN
a5ea8197ab
Anchor weather and other realtime questions on runner /v1/search results (or an explicit refuse constraint) so the model cannot invent dates or temperatures from training memory.
97 lines
2.2 KiB
Go
97 lines
2.2 KiB
Go
package agent
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
const piRunnerSearchTimeout = 20 * time.Second
|
|
|
|
// PiRunnerSearcher 通过 pi-agent runner 的 /v1/search 做 Tavily 检索。
|
|
type PiRunnerSearcher struct {
|
|
URL string
|
|
Token string
|
|
Client *http.Client
|
|
}
|
|
|
|
// Search 调用 runner /v1/search，返回 markdown 结果。
|
|
func (s *PiRunnerSearcher) Search(ctx context.Context, query string) (string, error) {
|
|
if s == nil {
|
|
return "", fmt.Errorf("pi runner searcher is nil")
|
|
}
|
|
base := strings.TrimRight(strings.TrimSpace(s.URL), "/")
|
|
q := strings.TrimSpace(query)
|
|
if base == "" {
|
|
return "", fmt.Errorf("pi runner url empty")
|
|
}
|
|
if q == "" {
|
|
return "", fmt.Errorf("search query empty")
|
|
}
|
|
|
|
body, err := json.Marshal(map[string]any{
|
|
"query": q,
|
|
"max_results": 5,
|
|
})
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
reqCtx, cancel := context.WithTimeout(ctx, piRunnerSearchTimeout)
|
|
defer cancel()
|
|
|
|
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, base+"/v1/search", bytes.NewReader(body))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
ApplyRunnerAuth(httpReq, s.Token)
|
|
|
|
client := s.Client
|
|
if client == nil {
|
|
client = &http.Client{Timeout: piRunnerSearchTimeout}
|
|
}
|
|
resp, err := client.Do(httpReq)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
limited, err := io.ReadAll(io.LimitReader(resp.Body, MaxResultBodyBytes+1))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if len(limited) > MaxResultBodyBytes {
|
|
return "", fmt.Errorf("search response too large")
|
|
}
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return "", fmt.Errorf("search http %d", resp.StatusCode)
|
|
}
|
|
|
|
var parsed struct {
|
|
OK bool `json:"ok"`
|
|
Markdown string `json:"markdown"`
|
|
Error string `json:"error"`
|
|
}
|
|
if err := json.Unmarshal(limited, &parsed); err != nil {
|
|
return "", fmt.Errorf("decode search response: %w", err)
|
|
}
|
|
if !parsed.OK {
|
|
msg := strings.TrimSpace(parsed.Error)
|
|
if msg == "" {
|
|
msg = "search failed"
|
|
}
|
|
return "", fmt.Errorf("%s", msg)
|
|
}
|
|
md := strings.TrimSpace(parsed.Markdown)
|
|
if md == "" {
|
|
return "", fmt.Errorf("empty search markdown")
|
|
}
|
|
return md, nil
|
|
}
|