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.
70 lines
1.9 KiB
Go
70 lines
1.9 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestPiRunnerSearcher_Success(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1/search" {
|
|
t.Fatalf("path=%s", r.URL.Path)
|
|
}
|
|
if r.Header.Get("Authorization") != "Bearer tok-1" {
|
|
t.Fatalf("missing auth: %q", r.Header.Get("Authorization"))
|
|
}
|
|
var body map[string]any
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
t.Fatalf("decode: %v", err)
|
|
}
|
|
if body["query"] != "今天北京天气" {
|
|
t.Fatalf("query=%v", body["query"])
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"ok": true,
|
|
"markdown": "## 北京天气\n气温 21~33°C",
|
|
})
|
|
}))
|
|
defer srv.Close()
|
|
|
|
s := &PiRunnerSearcher{URL: srv.URL, Token: "tok-1", Client: srv.Client()}
|
|
md, err := s.Search(context.Background(), "今天北京天气")
|
|
if err != nil {
|
|
t.Fatalf("Search: %v", err)
|
|
}
|
|
if !strings.Contains(md, "21~33") {
|
|
t.Fatalf("markdown=%q", md)
|
|
}
|
|
}
|
|
|
|
func TestPiRunnerSearcher_HTTPFailure(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusBadGateway)
|
|
_, _ = w.Write([]byte(`{"ok":false,"error":"upstream"}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
s := &PiRunnerSearcher{URL: srv.URL, Client: srv.Client()}
|
|
_, err := s.Search(context.Background(), "今天天气")
|
|
if err == nil {
|
|
t.Fatal("expected error")
|
|
}
|
|
}
|
|
|
|
func TestPiRunnerSearcher_EmptyMarkdown(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"ok": true, "markdown": " "})
|
|
}))
|
|
defer srv.Close()
|
|
|
|
s := &PiRunnerSearcher{URL: srv.URL, Client: srv.Client()}
|
|
_, err := s.Search(context.Background(), "新闻")
|
|
if err == nil {
|
|
t.Fatal("expected empty markdown error")
|
|
}
|
|
}
|