mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-28 03:38:11 +08:00
feat(agent): Go port — canvas engine, 22 components, DSL v2, 13 endpoints (#15952)
Ports the agent canvas subsystem from Python to Go.
## What's included
### Canvas Engine (Phase 0/1)
- State engine, scheduler, variable resolver, Redis checkpoint store,
cancel protocol
- **209 tests** across canvas / component / io packages
### 22 Components (P0–P4)
| Tier | Components |
|---|---|
| P0 T1+T2+T3 | LLM, Agent, ExitLoop, Switch, Categorize, Begin,
Message, Invoke |
| P1 T3 | VariableAggregator, VariableAssigner, StringTransform,
ListOperations, DataOperations |
| P2 T3 | Iteration, IterationItem, Loop, LoopItem |
| P3 T3 | UserFillUp, Fillup |
| P4 T5 | Browser, ExcelProcessor, DocsGenerator |
### DSL v2 Schema (Phase 2.5)
- Typed v2 in-memory model with v1-to-v2 auto-detect converter
- v1 legacy field stripping per plan §2.11.7
### HTTP Endpoints & Bug Fixes (Plans PR1–PR3)
- **DELETE SQL bug fix**: gorm v2 `Where("id = ?", id).Delete(...)`
pattern
- **CreateAgent validation**: title/DSL required, duplicate check, 103
envelope
- **13 new endpoints**: templates, prompts, tags, sessions CRUD,
chat/completions (SSE + non-stream stubs), rerun, test_db_connection,
logs, webhook/logs
- **756 Go unit tests** (745 → 756, +18)
- **17 → 0 Python integration test failures** (test_agents.py +
test_session_management/)
### Tools
21 eino tools: HTTPHelper, search tools, financial/data tools, mandatory
stubs
### Infrastructure
OTel observability, NATS message queue, DeepDoc gRPC client, SSRF
guards, IDOR mitigation
This commit is contained in:
124
internal/agent/tool/akshare.go
Normal file
124
internal/agent/tool/akshare.go
Normal file
@@ -0,0 +1,124 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const akshareToolName = "akshare"
|
||||
|
||||
const akshareToolDescription = "Query Chinese A-share financial data (AkShare). " +
|
||||
"STUB: AkShare is a Python library — not available in the Go Canvas. " +
|
||||
"Use the Python Canvas for Chinese A-share data."
|
||||
|
||||
const akshareUnsupportedMessage = "AkShare is a Python library — not available in Go Canvas. " +
|
||||
"Use Python Canvas for Chinese A-share data."
|
||||
|
||||
// akshareParams is the JSON shape the model sends into InvokableRun.
|
||||
//
|
||||
// AkShare exposes 200+ indicator functions (e.g. stock_zh_a_hist,
|
||||
// stock_zh_a_spot_em, bond_zh_hs_cov). For a future HTTP shim that routes
|
||||
// to a Python backend we keep the same shape: a stock/asset symbol plus
|
||||
// the indicator name. The Go stub validates the shape and rejects all
|
||||
// invocations with a clear error.
|
||||
type akshareParams struct {
|
||||
Symbol string `json:"symbol"`
|
||||
Indicator string `json:"indicator"`
|
||||
}
|
||||
|
||||
// akshareEnvelope is the model-facing JSON shape. The stub always
|
||||
// returns a populated Error so the model gets a deterministic, parseable
|
||||
// failure.
|
||||
type akshareEnvelope struct {
|
||||
Fields []string `json:"fields,omitempty"`
|
||||
Items []any `json:"items,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// AkShareTool is the Phase 3 batch 4 stub implementation of the AkShare
|
||||
// Chinese financial data tool (plan §2.11.4 row 1, §5 Phase 3 第 4 批).
|
||||
//
|
||||
// AkShare is a Python library (https://github.com/akfamily/akshare) —
|
||||
// there is no public HTTP API. A future thin shim that routes to a
|
||||
// Python backend is deferred. For P3-B4 the tool is registered with the
|
||||
// canvas so DSLs that reference "akshare" continue to parse, but every
|
||||
// invocation fails fast with a clear "use Python Canvas" message rather
|
||||
// than silently no-op'ing.
|
||||
//
|
||||
// AkShareTool does not own an HTTPHelper — it never makes network calls.
|
||||
type AkShareTool struct{}
|
||||
|
||||
// NewAkShareTool returns an AkShareTool. No HTTPHelper is allocated;
|
||||
// the stub never issues network requests.
|
||||
func NewAkShareTool() *AkShareTool { return &AkShareTool{} }
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (a *AkShareTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: akshareToolName,
|
||||
Desc: akshareToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"symbol": {
|
||||
Type: schema.String,
|
||||
Desc: "Stock or asset symbol (e.g. 000001, sh600000, bj920566).",
|
||||
Required: true,
|
||||
},
|
||||
"indicator": {
|
||||
Type: schema.String,
|
||||
Desc: "AkShare indicator function name (e.g. stock_zh_a_hist, stock_zh_a_spot_em).",
|
||||
Required: true,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// InvokableRun validates the input shape and returns a clear "use Python
|
||||
// Canvas" error. The model receives a JSON envelope with the message in
|
||||
// the `_ERROR` field so it can present a useful message back to the user
|
||||
// without surfacing a Go stack trace.
|
||||
func (a *AkShareTool) InvokableRun(_ context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p akshareParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return akshareErrJSON(errors.New(akshareUnsupportedMessage)),
|
||||
errors.New(akshareUnsupportedMessage)
|
||||
}
|
||||
if p.Symbol == "" || p.Indicator == "" {
|
||||
return akshareErrJSON(errors.New(akshareUnsupportedMessage)),
|
||||
errors.New(akshareUnsupportedMessage)
|
||||
}
|
||||
return akshareErrJSON(errors.New(akshareUnsupportedMessage)),
|
||||
errors.New(akshareUnsupportedMessage)
|
||||
}
|
||||
|
||||
func akshareJSON(env akshareEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"akshare: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func akshareErrJSON(err error) string {
|
||||
return akshareJSON(akshareEnvelope{Error: err.Error()})
|
||||
}
|
||||
110
internal/agent/tool/akshare_test.go
Normal file
110
internal/agent/tool/akshare_test.go
Normal file
@@ -0,0 +1,110 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAkShare_StubsUnsupported(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
args string
|
||||
wantSub string
|
||||
}{
|
||||
{
|
||||
name: "well-formed args",
|
||||
args: `{"symbol":"000001","indicator":"stock_zh_a_hist"}`,
|
||||
wantSub: "Python library",
|
||||
},
|
||||
{
|
||||
name: "missing indicator",
|
||||
args: `{"symbol":"000001"}`,
|
||||
wantSub: "Python library",
|
||||
},
|
||||
{
|
||||
name: "missing symbol",
|
||||
args: `{"indicator":"stock_zh_a_hist"}`,
|
||||
wantSub: "Python library",
|
||||
},
|
||||
{
|
||||
name: "empty payload",
|
||||
args: `{}`,
|
||||
wantSub: "Python library",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
tool := NewAkShareTool()
|
||||
out, err := tool.InvokableRun(context.Background(), tc.args)
|
||||
if err == nil {
|
||||
t.Fatalf("expected error, got nil (out=%s)", out)
|
||||
}
|
||||
if !strings.Contains(err.Error(), tc.wantSub) {
|
||||
t.Errorf("err = %q, want substring %q", err.Error(), tc.wantSub)
|
||||
}
|
||||
var env akshareEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error == "" {
|
||||
t.Errorf("env.Error = empty, want populated")
|
||||
}
|
||||
if !strings.Contains(env.Error, "Python") {
|
||||
t.Errorf("env.Error = %q, want to mention Python Canvas", env.Error)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAkShare_RejectsMalformedJSON(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewAkShareTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{not json`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for malformed JSON, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "Python") {
|
||||
t.Errorf("err = %q, want to mention Python", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAkShare_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewAkShareTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "akshare" {
|
||||
t.Errorf("Name = %q, want akshare", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "AkShare") {
|
||||
t.Errorf("Desc = %q, want to mention AkShare", info.Desc)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "STUB") && !strings.Contains(info.Desc, "Python") {
|
||||
t.Errorf("Desc = %q, want to flag stub status", info.Desc)
|
||||
}
|
||||
}
|
||||
253
internal/agent/tool/arxiv.go
Normal file
253
internal/agent/tool/arxiv.go
Normal file
@@ -0,0 +1,253 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"encoding/xml"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const arxivToolName = "arxiv"
|
||||
|
||||
const arxivToolDescription = "Search arXiv and return matching preprints as {title, authors, summary, pdf_url, entry_id}."
|
||||
|
||||
// arxivParams is the JSON shape the model sends into InvokableRun.
|
||||
type arxivParams struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
// arxivResult is one entry in the model-facing result list.
|
||||
type arxivResult struct {
|
||||
Title string `json:"title"`
|
||||
Authors []string `json:"authors"`
|
||||
Summary string `json:"summary"`
|
||||
PDFURL string `json:"pdf_url"`
|
||||
EntryID string `json:"entry_id"`
|
||||
}
|
||||
|
||||
// arxivEnvelope is the JSON shape the model sees.
|
||||
type arxivEnvelope struct {
|
||||
Results []arxivResult `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// arxivAtom is the upstream Atom 1.0 envelope. We only model the fields
|
||||
// we care about. Authors are stored as a flat list of <name> elements;
|
||||
// pdf_url is derived from the first <link> with rel="related" and
|
||||
// type="application/pdf", falling back to the entry id.
|
||||
type arxivAtom struct {
|
||||
XMLName xml.Name `xml:"http://www.w3.org/2005/Atom feed"`
|
||||
Title string `xml:"title"`
|
||||
Entries []arxivEntry `xml:"entry"`
|
||||
}
|
||||
|
||||
type arxivEntry struct {
|
||||
Title string `xml:"title"`
|
||||
ID string `xml:"id"`
|
||||
Summary string `xml:"summary"`
|
||||
Authors []arxivAuthor `xml:"author"`
|
||||
Links []arxivLink `xml:"link"`
|
||||
}
|
||||
|
||||
type arxivAuthor struct {
|
||||
Name string `xml:"name"`
|
||||
}
|
||||
|
||||
type arxivLink struct {
|
||||
Href string `xml:"href,attr"`
|
||||
Rel string `xml:"rel,attr"`
|
||||
Type string `xml:"type,attr"`
|
||||
}
|
||||
|
||||
// ArxivTool is the Phase 3 batch 2 implementation of the ArXiv academic
|
||||
// search tool (plan §2.11.4 row 2, §5 Phase 3 第 2 批). It performs a GET
|
||||
// against the public ArXiv API and parses the Atom XML response using
|
||||
// the stdlib encoding/xml package.
|
||||
type ArxivTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewArxivTool returns an ArxivTool using the default HTTPHelper.
|
||||
func NewArxivTool() *ArxivTool {
|
||||
return NewArxivToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewArxivToolWith returns an ArxivTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewArxivToolWith(h *HTTPHelper) *ArxivTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &ArxivTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (a *ArxivTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: arxivToolName,
|
||||
Desc: arxivToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query (matches the arXiv `all:` field).",
|
||||
Required: true,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of results to return. Defaults to 5.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildArxivURL constructs the ArXiv /api/query URL.
|
||||
func buildArxivURL(query string, maxResults int) string {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
q := url.Values{}
|
||||
q.Set("search_query", "all:"+query)
|
||||
q.Set("max_results", fmt.Sprintf("%d", maxResults))
|
||||
return "http://export.arxiv.org/api/query?" + q.Encode()
|
||||
}
|
||||
|
||||
// parseArxivAtom decodes the upstream Atom XML and returns the model-
|
||||
// facing results. Exposed at package scope for unit testing with canned
|
||||
// XML.
|
||||
func parseArxivAtom(body []byte) ([]arxivResult, error) {
|
||||
var feed arxivAtom
|
||||
if err := xml.NewDecoder(bytes.NewReader(body)).Decode(&feed); err != nil {
|
||||
return nil, fmt.Errorf("arxiv: decode atom: %w", err)
|
||||
}
|
||||
results := make([]arxivResult, 0, len(feed.Entries))
|
||||
for _, e := range feed.Entries {
|
||||
authors := make([]string, 0, len(e.Authors))
|
||||
for _, a := range e.Authors {
|
||||
if a.Name != "" {
|
||||
authors = append(authors, a.Name)
|
||||
}
|
||||
}
|
||||
pdfURL := pickArxivPDF(e)
|
||||
results = append(results, arxivResult{
|
||||
Title: normalizeArxivWhitespace(e.Title),
|
||||
Authors: authors,
|
||||
Summary: normalizeArxivWhitespace(e.Summary),
|
||||
PDFURL: pdfURL,
|
||||
EntryID: e.ID,
|
||||
})
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// pickArxivPDF returns the entry's PDF link, preferring the explicit
|
||||
// rel="related" type="application/pdf" link and falling back to the
|
||||
// canonical abs/.../pdf URL derived from the entry id.
|
||||
func pickArxivPDF(e arxivEntry) string {
|
||||
for _, l := range e.Links {
|
||||
if l.Rel == "related" && l.Type == "application/pdf" && l.Href != "" {
|
||||
return l.Href
|
||||
}
|
||||
}
|
||||
// Fallback: derive from the abs id. The arXiv API returns ids like
|
||||
// "http://arxiv.org/abs/2501.12345v1"; we transform the abs path
|
||||
// to a /pdf/ path.
|
||||
id := e.ID
|
||||
if id == "" {
|
||||
return ""
|
||||
}
|
||||
u, err := url.Parse(id)
|
||||
if err != nil || u.Path == "" {
|
||||
return id
|
||||
}
|
||||
newPath := u.Path
|
||||
if i := strings.LastIndex(newPath, "abs"); i >= 0 {
|
||||
newPath = newPath[:i] + "pdf" + newPath[i+len("abs"):]
|
||||
} else {
|
||||
newPath = "/pdf" + newPath
|
||||
}
|
||||
u2 := *u
|
||||
u2.Path = newPath
|
||||
return u2.String()
|
||||
}
|
||||
|
||||
// normalizeArxivWhitespace collapses internal whitespace and trims.
|
||||
// The arXiv XML wraps long fields across multiple <p> blocks, leaving
|
||||
// spurious newlines and double spaces when concatenated.
|
||||
func normalizeArxivWhitespace(s string) string {
|
||||
return strings.TrimSpace(strings.Join(strings.Fields(s), " "))
|
||||
}
|
||||
|
||||
// InvokableRun performs the ArXiv search.
|
||||
func (a *ArxivTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p arxivParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return arxivErrJSON(fmt.Errorf("arxiv: parse arguments: %w", err)),
|
||||
fmt.Errorf("arxiv: parse arguments: %w", err)
|
||||
}
|
||||
if p.Query == "" {
|
||||
return arxivErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("arxiv: query is required")
|
||||
}
|
||||
|
||||
endpoint := buildArxivURL(p.Query, p.MaxResults)
|
||||
resp, err := a.helper.Do(ctx, http.MethodGet, endpoint, "", "", nil)
|
||||
if err != nil {
|
||||
return arxivErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return arxivErrJSON(fmt.Errorf("arxiv: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("arxiv: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return arxivErrJSON(fmt.Errorf("arxiv: read body: %w", err)),
|
||||
fmt.Errorf("arxiv: read body: %w", err)
|
||||
}
|
||||
|
||||
results, err := parseArxivAtom(body)
|
||||
if err != nil {
|
||||
return arxivErrJSON(err), err
|
||||
}
|
||||
return arxivJSON(arxivEnvelope{Results: results}), nil
|
||||
}
|
||||
|
||||
func arxivJSON(env arxivEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"arxiv: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func arxivErrJSON(err error) string {
|
||||
return arxivJSON(arxivEnvelope{Error: err.Error()})
|
||||
}
|
||||
200
internal/agent/tool/arxiv_test.go
Normal file
200
internal/agent/tool/arxiv_test.go
Normal file
@@ -0,0 +1,200 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestArxiv_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got := buildArxivURL("transformer", 3)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != "export.arxiv.org" {
|
||||
t.Errorf("host = %q, want export.arxiv.org", u.Host)
|
||||
}
|
||||
if u.Path != "/api/query" {
|
||||
t.Errorf("path = %q, want /api/query", u.Path)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("search_query") != "all:transformer" {
|
||||
t.Errorf("search_query = %q, want all:transformer", q.Get("search_query"))
|
||||
}
|
||||
if q.Get("max_results") != "3" {
|
||||
t.Errorf("max_results = %q, want 3", q.Get("max_results"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestArxiv_ParseAtomEntry(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const canned = `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<feed xmlns="http://www.w3.org/2005/Atom" xmlns:arxiv="http://arxiv.org/schemas/atom">
|
||||
<title>arXiv Query: all:rag</title>
|
||||
<entry>
|
||||
<id>http://arxiv.org/abs/2501.12345v1</id>
|
||||
<title>Retrieval-Augmented Generation for
|
||||
Open-Domain Question Answering</title>
|
||||
<summary> We present a method
|
||||
combining retrieval and generation. </summary>
|
||||
<author><name>Alice Liddell</name></author>
|
||||
<author><name>Bob Builder</name></author>
|
||||
<link href="http://arxiv.org/abs/2501.12345v1" rel="alternate" type="text/html"/>
|
||||
<link href="http://arxiv.org/pdf/2501.12345v1" rel="related" type="application/pdf"/>
|
||||
</entry>
|
||||
<entry>
|
||||
<id>http://arxiv.org/abs/2409.99999v2</id>
|
||||
<title>Single-Author Paper</title>
|
||||
<summary>Brief.</summary>
|
||||
<author><name>Carol Danvers</name></author>
|
||||
<!-- no pdf link: pickArxivPDF must fall back to abs id derived pdf -->
|
||||
</entry>
|
||||
</feed>`
|
||||
|
||||
results, err := parseArxivAtom([]byte(canned))
|
||||
if err != nil {
|
||||
t.Fatalf("parseArxivAtom: %v", err)
|
||||
}
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("len(results) = %d, want 2", len(results))
|
||||
}
|
||||
|
||||
// First entry: explicit pdf link
|
||||
r0 := results[0]
|
||||
if !strings.Contains(r0.Title, "Retrieval-Augmented Generation") {
|
||||
t.Errorf("r0.Title = %q, want to contain 'Retrieval-Augmented Generation'", r0.Title)
|
||||
}
|
||||
if !strings.HasPrefix(r0.Title, "Retrieval-") {
|
||||
// whitespace normalized
|
||||
if strings.Contains(r0.Title, "\n") {
|
||||
t.Errorf("r0.Title contains raw newline: %q", r0.Title)
|
||||
}
|
||||
}
|
||||
if len(r0.Authors) != 2 {
|
||||
t.Errorf("r0.Authors len = %d, want 2", len(r0.Authors))
|
||||
}
|
||||
if r0.Authors[0] != "Alice Liddell" || r0.Authors[1] != "Bob Builder" {
|
||||
t.Errorf("r0.Authors = %v, want [Alice Liddell Bob Builder]", r0.Authors)
|
||||
}
|
||||
if r0.PDFURL != "http://arxiv.org/pdf/2501.12345v1" {
|
||||
t.Errorf("r0.PDFURL = %q, want http://arxiv.org/pdf/2501.12345v1", r0.PDFURL)
|
||||
}
|
||||
if r0.EntryID != "http://arxiv.org/abs/2501.12345v1" {
|
||||
t.Errorf("r0.EntryID = %q, want http://arxiv.org/abs/2501.12345v1", r0.EntryID)
|
||||
}
|
||||
if !strings.HasPrefix(r0.Summary, "We present") {
|
||||
t.Errorf("r0.Summary = %q, want to start with 'We present' (whitespace normalized)", r0.Summary)
|
||||
}
|
||||
|
||||
// Second entry: fallback pdf derived from abs id
|
||||
r1 := results[1]
|
||||
if r1.PDFURL != "http://arxiv.org/pdf/2409.99999v2" {
|
||||
t.Errorf("r1.PDFURL = %q, want http://arxiv.org/pdf/2409.99999v2 (derived)", r1.PDFURL)
|
||||
}
|
||||
if len(r1.Authors) != 1 || r1.Authors[0] != "Carol Danvers" {
|
||||
t.Errorf("r1.Authors = %v, want [Carol Danvers]", r1.Authors)
|
||||
}
|
||||
}
|
||||
|
||||
func TestArxiv_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewArxivTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "arxiv" {
|
||||
t.Errorf("Name = %q, want arxiv", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "arXiv") {
|
||||
t.Errorf("Desc = %q, want to mention arXiv", info.Desc)
|
||||
}
|
||||
}
|
||||
|
||||
func TestArxiv_RequiresQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewArxivTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{"query":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty query")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "query") {
|
||||
t.Errorf("err = %v, want to mention query", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestArxiv_FullRoundtrip(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const canned = `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<feed xmlns="http://www.w3.org/2005/Atom">
|
||||
<title>arXiv Query</title>
|
||||
<entry>
|
||||
<id>http://arxiv.org/abs/2501.12345v1</id>
|
||||
<title>Test Paper</title>
|
||||
<summary>Summary.</summary>
|
||||
<author><name>Author One</name></author>
|
||||
<link href="http://arxiv.org/pdf/2501.12345v1" rel="related" type="application/pdf"/>
|
||||
</entry>
|
||||
</feed>`
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/atom+xml")
|
||||
_, _ = w.Write([]byte(canned))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
// rewriteHostTransport points the hard-coded export.arxiv.org at the
|
||||
// test server.
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewArxivToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(), `{"query":"rag","max_results":5}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env arxivEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 1 {
|
||||
t.Fatalf("Results len = %d, want 1", len(env.Results))
|
||||
}
|
||||
if env.Results[0].Title != "Test Paper" {
|
||||
t.Errorf("Title = %q, want Test Paper", env.Results[0].Title)
|
||||
}
|
||||
if env.Results[0].PDFURL != "http://arxiv.org/pdf/2501.12345v1" {
|
||||
t.Errorf("PDFURL = %q, want http://arxiv.org/pdf/2501.12345v1", env.Results[0].PDFURL)
|
||||
}
|
||||
}
|
||||
158
internal/agent/tool/code_exec.go
Normal file
158
internal/agent/tool/code_exec.go
Normal file
@@ -0,0 +1,158 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// ErrCodeExecSandboxMissing is returned when the Phase 5 gRPC client to the
|
||||
// Python sandbox (agent.sandbox.client.execute_code) is not yet wired. The
|
||||
// Python sandbox itself is kept as-is per plan §2.11.4 (do NOT rewrite the
|
||||
// sandbox); Phase 3 batch 1 only ships the tool shell, and Phase 5 adds
|
||||
// the gRPC client.
|
||||
var ErrCodeExecSandboxMissing = errors.New(
|
||||
"CodeExec sandbox gRPC client not yet implemented in Go — " +
|
||||
"defer to Python Canvas",
|
||||
)
|
||||
|
||||
const codeExecToolName = "execute_code"
|
||||
|
||||
const codeExecToolDescription = "This tool has a sandbox that can execute code written in 'Python'/'Javascript'. " +
|
||||
"It receives a piece of code and returns a JSON string."
|
||||
|
||||
// codeExecArgs is the JSON shape the model sends in. The Python tool
|
||||
// accepts "lang" + "script" (see agent/tools/code_exec.py:303-310); we
|
||||
// also accept "code" as a synonym since some DSLs and Phase 3 batch 1
|
||||
// tests use that spelling.
|
||||
type codeExecArgs struct {
|
||||
Language string `json:"language,omitempty"`
|
||||
Lang string `json:"lang,omitempty"`
|
||||
Script string `json:"script,omitempty"`
|
||||
Code string `json:"code,omitempty"`
|
||||
Args map[string]any `json:"arguments,omitempty"`
|
||||
}
|
||||
|
||||
// codeExecResult is the JSON envelope returned to the model. The output
|
||||
// shape mirrors the Python tool's `content` / `_ERROR` / `actual_type`
|
||||
// fields so downstream nodes can pattern-match unchanged.
|
||||
type codeExecResult struct {
|
||||
Content string `json:"content,omitempty"`
|
||||
ActualType string `json:"actual_type,omitempty"`
|
||||
Stub bool `json:"stub,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// CodeExecTool is the Phase 3 batch 1 shell for the CodeExec tool
|
||||
// (plan §2.11.4 row 3, §5 Phase 3 第 1 批). It validates language +
|
||||
// non-empty code and returns a structured "not-yet-wired" error.
|
||||
type CodeExecTool struct{}
|
||||
|
||||
// NewCodeExecTool returns a CodeExecTool implementing eino's
|
||||
// tool.InvokableTool interface.
|
||||
func NewCodeExecTool() *CodeExecTool {
|
||||
return &CodeExecTool{}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model. The schema mirrors
|
||||
// the Python CodeExecParam ToolMeta (plan §5 Phase 3, 字段对齐).
|
||||
func (c *CodeExecTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: codeExecToolName,
|
||||
Desc: codeExecToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"language": {
|
||||
Type: schema.String,
|
||||
Desc: "The programming language of the code. Allowed: 'python' (or 'python3'), 'javascript' (or 'nodejs').",
|
||||
Enum: []string{"python", "python3", "javascript", "nodejs"},
|
||||
Required: true,
|
||||
},
|
||||
"code": {
|
||||
Type: schema.String,
|
||||
Desc: "The code to execute. Must define a `main` function (Python) or export `main` (JavaScript).",
|
||||
Required: true,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// InvokableRun validates the inputs and returns the structured stub error.
|
||||
// Phase 5 will replace the stub body with a gRPC call into the Python
|
||||
// sandbox (per plan §2.11.4 "不重写沙箱" 决策).
|
||||
func (c *CodeExecTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
|
||||
var args codeExecArgs
|
||||
if argumentsInJSON == "" {
|
||||
return codeExecStubResult("arguments are required"), errors.New("code_exec: empty arguments")
|
||||
}
|
||||
if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil {
|
||||
return codeExecStubResult("invalid JSON: " + err.Error()),
|
||||
fmt.Errorf("code_exec: parse arguments: %w", err)
|
||||
}
|
||||
|
||||
lang := normalizeCodeExecLang(args.Language, args.Lang)
|
||||
if lang == "" {
|
||||
return codeExecStubResult("unsupported language: must be 'python'/'python3'/'javascript'/'nodejs'"),
|
||||
errors.New("code_exec: invalid language")
|
||||
}
|
||||
|
||||
script := args.Script
|
||||
if script == "" {
|
||||
script = args.Code
|
||||
}
|
||||
if strings.TrimSpace(script) == "" {
|
||||
return codeExecStubResult("code is required"), errors.New("code_exec: empty code")
|
||||
}
|
||||
|
||||
// Phase 3 batch 1: gRPC client wires in Phase 5.
|
||||
return codeExecStubResult(ErrCodeExecSandboxMissing.Error()),
|
||||
ErrCodeExecSandboxMissing
|
||||
}
|
||||
|
||||
func codeExecStubResult(msg string) string {
|
||||
b, err := json.Marshal(codeExecResult{
|
||||
Stub: true,
|
||||
Error: msg,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"code_exec: marshal stub: %s","stub":true}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// normalizeCodeExecLang accepts the model's literal "language" or the
|
||||
// Python-style "lang" alias and maps synonyms to the canonical "python" /
|
||||
// "nodejs" forms used by the Python sandbox.
|
||||
func normalizeCodeExecLang(primary, alias string) string {
|
||||
v := strings.ToLower(strings.TrimSpace(primary))
|
||||
if v == "" {
|
||||
v = strings.ToLower(strings.TrimSpace(alias))
|
||||
}
|
||||
switch v {
|
||||
case "python", "python3":
|
||||
return "python"
|
||||
case "javascript", "js", "nodejs", "node":
|
||||
return "nodejs"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
94
internal/agent/tool/code_exec_test.go
Normal file
94
internal/agent/tool/code_exec_test.go
Normal file
@@ -0,0 +1,94 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCodeExec_StubsErrorWhenClientMissing(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCodeExecTool()
|
||||
out, err := c.InvokableRun(context.Background(), `{"language":"python","code":"def main(): return {}"}`)
|
||||
if !errors.Is(err, ErrCodeExecSandboxMissing) {
|
||||
t.Fatalf("err = %v, want ErrCodeExecSandboxMissing", err)
|
||||
}
|
||||
|
||||
var got codeExecResult
|
||||
if jerr := json.Unmarshal([]byte(out), &got); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if !got.Stub {
|
||||
t.Errorf("Stub = false, want true")
|
||||
}
|
||||
if !strings.Contains(got.Error, "sandbox") {
|
||||
t.Errorf("Error = %q, want to mention 'sandbox'", got.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodeExec_RejectsEmptyCode(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCodeExecTool()
|
||||
_, err := c.InvokableRun(context.Background(), `{"language":"python","code":""}`)
|
||||
if err == nil || !strings.Contains(err.Error(), "code") {
|
||||
t.Fatalf("err = %v, want to mention empty code", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodeExec_RejectsBadLanguage(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCodeExecTool()
|
||||
_, err := c.InvokableRun(context.Background(), `{"language":"brainfuck","code":"x"}`)
|
||||
if err == nil || !strings.Contains(err.Error(), "language") {
|
||||
t.Fatalf("err = %v, want to reject unsupported language", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodeExec_AcceptsLangAlias(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCodeExecTool()
|
||||
// Python tool also accepts "lang" as the field name; the Go shell
|
||||
// should still reach the stub branch.
|
||||
_, err := c.InvokableRun(context.Background(), `{"lang":"nodejs","script":"async function main() {}"}`)
|
||||
if !errors.Is(err, ErrCodeExecSandboxMissing) {
|
||||
t.Fatalf("err = %v, want ErrCodeExecSandboxMissing", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodeExec_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCodeExecTool()
|
||||
info, err := c.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "execute_code" {
|
||||
t.Errorf("Name = %q, want execute_code", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Python") {
|
||||
t.Errorf("Desc = %q, want to mention Python", info.Desc)
|
||||
}
|
||||
}
|
||||
306
internal/agent/tool/crawler.go
Normal file
306
internal/agent/tool/crawler.go
Normal file
@@ -0,0 +1,306 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"golang.org/x/net/html"
|
||||
)
|
||||
|
||||
const crawlerToolName = "crawler"
|
||||
|
||||
const crawlerToolDescription = "Fetches a web page and returns its extracted text content and links."
|
||||
|
||||
// crawlerArgs is the JSON shape the model sends in. max_depth and
|
||||
// max_pages are accepted for API symmetry with the Python tool, but
|
||||
// Phase 3 batch 1 only implements depth=0 (single page fetch).
|
||||
type crawlerArgs struct {
|
||||
URL string `json:"url"`
|
||||
MaxDepth int `json:"max_depth,omitempty"`
|
||||
MaxPages int `json:"max_pages,omitempty"`
|
||||
}
|
||||
|
||||
// crawlerResult is the JSON envelope returned to the model. The shape
|
||||
// mirrors the Python tool's `content` / `links` output.
|
||||
type crawlerResult struct {
|
||||
URL string `json:"url"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
Links []string `json:"links,omitempty"`
|
||||
Status int `json:"status,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// Resolver validates a URL and returns the pinned IP for the host. The
|
||||
// returned IP is dialed directly by HTTPHelper.DoPinned which defeats
|
||||
// DNS rebinding: an attacker cannot swap a public record for a private
|
||||
// one between the resolver's lookup and the actual connect, because
|
||||
// the connect is pinned at the *http.Transport dialer layer (see
|
||||
// pinnedDialer in http_helper.go) and never re-resolves the hostname.
|
||||
// The request URL host is preserved so TLS SNI and cert verification
|
||||
// continue to target the validated hostname.
|
||||
//
|
||||
// The default production resolver is ResolveAndValidate (ssrf.go),
|
||||
// which rejects loopback / link-local / private / metadata targets and
|
||||
// returns the first safe A/AAAA record.
|
||||
type Resolver func(rawURL string) (host string, ip net.IP, err error)
|
||||
|
||||
// CrawlerTool is the Phase 3 batch 1 implementation of the Crawler tool
|
||||
// (plan §2.11.4 row 4, §5 Phase 3 第 1 批). It fetches a single page
|
||||
// (max_depth=0) via HTTPHelper and extracts text + links with
|
||||
// golang.org/x/net/html.
|
||||
type CrawlerTool struct {
|
||||
helper *HTTPHelper
|
||||
// resolve is the URL resolver used to block internal / metadata
|
||||
// targets AND to pin the host to a known-safe IP. It is a function
|
||||
// field (rather than a hard call to ResolveAndValidate) so unit tests
|
||||
// that use httptest.NewServer (which binds to 127.0.0.1) can swap in
|
||||
// a no-op that returns the literal IP. Production construction
|
||||
// always uses ResolveAndValidate.
|
||||
resolve Resolver
|
||||
}
|
||||
|
||||
// NewCrawlerTool returns a CrawlerTool using the default HTTPHelper.
|
||||
// Pass NewCrawlerToolWith(helper) to inject a custom HTTPHelper (e.g.
|
||||
// with a test transport).
|
||||
func NewCrawlerTool() *CrawlerTool {
|
||||
return NewCrawlerToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewCrawlerToolWith returns a CrawlerTool that uses the provided
|
||||
// HTTPHelper. Useful for tests and for sharing a single helper across
|
||||
// multiple tool instances.
|
||||
func NewCrawlerToolWith(h *HTTPHelper) *CrawlerTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &CrawlerTool{helper: h, resolve: ResolveAndValidate}
|
||||
}
|
||||
|
||||
// WithResolver replaces the URL resolver (which performs the SSRF
|
||||
// check and supplies the pinned IP) with a custom function. The default
|
||||
// is ResolveAndValidate; tests that point the crawler at an
|
||||
// httptest.NewServer (127.0.0.1) can pass a no-op that returns the
|
||||
// literal host. Returns the same receiver for fluent use.
|
||||
func (c *CrawlerTool) WithResolver(fn Resolver) *CrawlerTool {
|
||||
if fn != nil {
|
||||
c.resolve = fn
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (c *CrawlerTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: crawlerToolName,
|
||||
Desc: crawlerToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"url": {
|
||||
Type: schema.String,
|
||||
Desc: "The URL to fetch. Must be http or https.",
|
||||
Required: true,
|
||||
},
|
||||
"max_depth": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Recursion depth. Phase 3 batch 1 supports 0 only; >0 returns an error.",
|
||||
Required: false,
|
||||
},
|
||||
"max_pages": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of pages to fetch. Phase 3 batch 1 ignores this (single page).",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ErrCrawlerDepthUnsupported is returned when the caller asks for
|
||||
// max_depth>0. Multi-page crawling is out of scope for Phase 3 batch 1.
|
||||
var ErrCrawlerDepthUnsupported = errors.New(
|
||||
"crawler: max_depth > 0 is not supported in Phase 3 batch 1; " +
|
||||
"use a single-page fetch (max_depth=0)",
|
||||
)
|
||||
|
||||
// InvokableRun fetches a single page and returns extracted text + links.
|
||||
// max_depth>0 is rejected; multi-page crawling is deferred to a later
|
||||
// batch.
|
||||
func (c *CrawlerTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
|
||||
var args crawlerArgs
|
||||
if argumentsInJSON == "" {
|
||||
return crawlerStubResult(crawlerResult{Error: "arguments are required"}),
|
||||
errors.New("crawler: empty arguments")
|
||||
}
|
||||
if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil {
|
||||
return crawlerStubResult(crawlerResult{Error: "invalid JSON: " + err.Error()}),
|
||||
fmt.Errorf("crawler: parse arguments: %w", err)
|
||||
}
|
||||
|
||||
if strings.TrimSpace(args.URL) == "" {
|
||||
return crawlerStubResult(crawlerResult{Error: "url is required"}),
|
||||
errors.New("crawler: empty url")
|
||||
}
|
||||
if !strings.HasPrefix(args.URL, "http://") && !strings.HasPrefix(args.URL, "https://") {
|
||||
return crawlerStubResult(crawlerResult{URL: args.URL, Error: "url must be http or https"}),
|
||||
fmt.Errorf("crawler: unsupported url scheme: %s", args.URL)
|
||||
}
|
||||
// Reject max_depth > 0 BEFORE the SSRF guard: the guard performs a
|
||||
// DNS lookup that may be slow / fail in CI, and a depth-0 caller
|
||||
// asking for max_depth=10 should be rejected on a structural
|
||||
// problem first.
|
||||
if args.MaxDepth > 0 {
|
||||
return crawlerStubResult(crawlerResult{URL: args.URL, Error: ErrCrawlerDepthUnsupported.Error()}),
|
||||
ErrCrawlerDepthUnsupported
|
||||
}
|
||||
// SSRF guard + DNS-rebinding pinning. c.resolve validates the URL
|
||||
// and returns the IP we should dial directly. DoPinned installs a
|
||||
// transport-level pinned dialer that connects to that IP, while the
|
||||
// request URL host stays as the original hostname — so an attacker
|
||||
// who flips the A record to a private address after this point
|
||||
// still cannot redirect the request (the connect is pinned to the
|
||||
// IP we resolved here) AND TLS SNI / cert verification continue to
|
||||
// target the validated hostname. Rewriting the URL host to the IP
|
||||
// would have broken HTTPS, so the pinning happens in the dialer.
|
||||
host, pinnedIP, resolveErr := c.resolve(args.URL)
|
||||
if resolveErr != nil {
|
||||
return crawlerStubResult(crawlerResult{URL: args.URL, Error: resolveErr.Error()}), resolveErr
|
||||
}
|
||||
|
||||
resp, err := c.helper.DoPinned(ctx, http.MethodGet, args.URL, "", "", nil, host, pinnedIP)
|
||||
if err != nil {
|
||||
return crawlerStubResult(crawlerResult{URL: args.URL, Error: err.Error()}), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return crawlerStubResult(crawlerResult{URL: args.URL, Status: resp.StatusCode, Error: "read body: " + err.Error()}),
|
||||
fmt.Errorf("crawler: read body: %w", err)
|
||||
}
|
||||
|
||||
page, err := extractPage(body)
|
||||
if err != nil {
|
||||
return crawlerStubResult(crawlerResult{URL: args.URL, Status: resp.StatusCode, Error: err.Error()}),
|
||||
fmt.Errorf("crawler: extract: %w", err)
|
||||
}
|
||||
page.URL = args.URL
|
||||
page.Status = resp.StatusCode
|
||||
|
||||
return crawlerJSON(page)
|
||||
}
|
||||
|
||||
// extractPage parses the HTML body and returns its title, plain text
|
||||
// content, and absolute links. It uses golang.org/x/net/html per plan
|
||||
// §2.11.4 (T2: HTTP + golang.org/x/net/html).
|
||||
func extractPage(body []byte) (crawlerResult, error) {
|
||||
doc, err := html.Parse(strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return crawlerResult{}, fmt.Errorf("parse html: %w", err)
|
||||
}
|
||||
|
||||
var out crawlerResult
|
||||
var titleNodes []*html.Node
|
||||
var text strings.Builder
|
||||
|
||||
var walk func(n *html.Node)
|
||||
walk = func(n *html.Node) {
|
||||
if n.Type == html.ElementNode {
|
||||
switch strings.ToLower(n.Data) {
|
||||
case "script", "style", "noscript", "template", "svg":
|
||||
// skip non-content subtrees entirely
|
||||
return
|
||||
case "head":
|
||||
// recurse into head only to capture <title>; skip text
|
||||
// (so meta tags etc. don't pollute the body text)
|
||||
for c := n.FirstChild; c != nil; c = c.NextSibling {
|
||||
if c.Type == html.ElementNode && strings.EqualFold(c.Data, "title") {
|
||||
titleNodes = append(titleNodes, c)
|
||||
}
|
||||
}
|
||||
return
|
||||
case "title":
|
||||
titleNodes = append(titleNodes, n)
|
||||
case "a":
|
||||
for _, a := range n.Attr {
|
||||
if strings.EqualFold(a.Key, "href") {
|
||||
href := strings.TrimSpace(a.Val)
|
||||
if href != "" && !strings.HasPrefix(href, "#") {
|
||||
out.Links = append(out.Links, href)
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if n.Type == html.TextNode {
|
||||
t := strings.TrimSpace(n.Data)
|
||||
if t != "" {
|
||||
if text.Len() > 0 {
|
||||
text.WriteByte(' ')
|
||||
}
|
||||
text.WriteString(t)
|
||||
}
|
||||
}
|
||||
for c := n.FirstChild; c != nil; c = c.NextSibling {
|
||||
walk(c)
|
||||
}
|
||||
}
|
||||
walk(doc)
|
||||
|
||||
// Extract title text (concatenate text nodes inside <title>).
|
||||
for _, tn := range titleNodes {
|
||||
var t strings.Builder
|
||||
for c := tn.FirstChild; c != nil; c = c.NextSibling {
|
||||
if c.Type == html.TextNode {
|
||||
t.WriteString(c.Data)
|
||||
}
|
||||
}
|
||||
title := strings.TrimSpace(t.String())
|
||||
if title != "" {
|
||||
out.Title = title
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
out.Content = strings.Join(strings.Fields(text.String()), " ")
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func crawlerStubResult(r crawlerResult) string {
|
||||
b, err := json.Marshal(r)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"crawler: marshal: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func crawlerJSON(r crawlerResult) (string, error) {
|
||||
b, err := json.Marshal(r)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("crawler: marshal result: %w", err)
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
152
internal/agent/tool/crawler_test.go
Normal file
152
internal/agent/tool/crawler_test.go
Normal file
@@ -0,0 +1,152 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const sampleHTML = `<!DOCTYPE html>
|
||||
<html>
|
||||
<head><title>RAGFlow Docs</title></head>
|
||||
<body>
|
||||
<h1>Welcome</h1>
|
||||
<p>This is a paragraph with <a href="/docs">docs link</a> and <a href="https://example.com">external</a>.</p>
|
||||
<script>var secret = "ignored";</script>
|
||||
<style>body { color: red; }</style>
|
||||
</body>
|
||||
</html>`
|
||||
|
||||
func TestCrawler_FetchesAndExtractsText(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
_, _ = w.Write([]byte(sampleHTML))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
// httptest.NewServer binds to 127.0.0.1 which the production
|
||||
// SSRF guard correctly blocks. Bypass it for this test only by
|
||||
// installing a resolver that returns the literal host/IP. This
|
||||
// mirrors the previous WithSSRFValidator behaviour and keeps the
|
||||
// test exercising the pinned-connect path (DoPinned), which is the
|
||||
// whole point of the M1-rebinding fix.
|
||||
loopbackResolver := func(rawURL string) (string, net.IP, error) {
|
||||
u, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
host := u.Hostname()
|
||||
return host, net.ParseIP(host), nil
|
||||
}
|
||||
c := NewCrawlerTool().WithResolver(loopbackResolver)
|
||||
out, err := c.InvokableRun(context.Background(),
|
||||
`{"url":`+jsonString(srv.URL)+`,"max_depth":0}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var got crawlerResult
|
||||
if jerr := json.Unmarshal([]byte(out), &got); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if got.Status != http.StatusOK {
|
||||
t.Errorf("Status = %d, want 200", got.Status)
|
||||
}
|
||||
if !strings.Contains(got.Title, "RAGFlow Docs") {
|
||||
t.Errorf("Title = %q, want to contain 'RAGFlow Docs'", got.Title)
|
||||
}
|
||||
if !strings.Contains(got.Content, "Welcome") {
|
||||
t.Errorf("Content = %q, want to contain 'Welcome'", got.Content)
|
||||
}
|
||||
// <script> and <style> contents must be stripped.
|
||||
if strings.Contains(got.Content, "secret") {
|
||||
t.Errorf("Content leaked <script> text: %q", got.Content)
|
||||
}
|
||||
if strings.Contains(got.Content, "color: red") {
|
||||
t.Errorf("Content leaked <style> text: %q", got.Content)
|
||||
}
|
||||
// Both links should be captured.
|
||||
wantLinks := map[string]bool{"/docs": false, "https://example.com": false}
|
||||
for _, l := range got.Links {
|
||||
if _, ok := wantLinks[l]; ok {
|
||||
wantLinks[l] = true
|
||||
}
|
||||
}
|
||||
for href, seen := range wantLinks {
|
||||
if !seen {
|
||||
t.Errorf("missing link %q in result", href)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCrawler_RejectsMaxDepthGreaterThanZero(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCrawlerTool()
|
||||
_, err := c.InvokableRun(context.Background(), `{"url":"https://example.com","max_depth":1}`)
|
||||
if !errors.Is(err, ErrCrawlerDepthUnsupported) {
|
||||
t.Fatalf("err = %v, want ErrCrawlerDepthUnsupported", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCrawler_RejectsMissingURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCrawlerTool()
|
||||
_, err := c.InvokableRun(context.Background(), `{"url":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty url")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCrawler_RejectsNonHTTPScheme(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCrawlerTool()
|
||||
_, err := c.InvokableRun(context.Background(), `{"url":"file:///etc/passwd"}`)
|
||||
if err == nil || !strings.Contains(err.Error(), "scheme") {
|
||||
t.Fatalf("err = %v, want to reject file:// scheme", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCrawler_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := NewCrawlerTool()
|
||||
info, err := c.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "crawler" {
|
||||
t.Errorf("Name = %q, want crawler", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "text") {
|
||||
t.Errorf("Desc = %q, want to mention text extraction", info.Desc)
|
||||
}
|
||||
}
|
||||
|
||||
// jsonString is defined in exesql_test.go (same package).
|
||||
198
internal/agent/tool/deepl.go
Normal file
198
internal/agent/tool/deepl.go
Normal file
@@ -0,0 +1,198 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const deeplToolName = "deepl"
|
||||
|
||||
const deeplToolDescription = "Translate text via the DeepL API. Returns translations[].{text, detected_source_language}."
|
||||
|
||||
// deeplParams is the JSON shape the model sends into InvokableRun.
|
||||
type deeplParams struct {
|
||||
APIKey string `json:"api_key"`
|
||||
Text string `json:"text"`
|
||||
SourceLang string `json:"source_lang"`
|
||||
TargetLang string `json:"target_lang"`
|
||||
}
|
||||
|
||||
// deeplTranslation is one element of the upstream `translations` array.
|
||||
type deeplTranslation struct {
|
||||
Text string `json:"text"`
|
||||
DetectedSourceLanguage string `json:"detected_source_language"`
|
||||
}
|
||||
|
||||
// deeplResponse is the upstream DeepL envelope.
|
||||
type deeplResponse struct {
|
||||
Translations []deeplTranslation `json:"translations"`
|
||||
}
|
||||
|
||||
// deeplEnvelope is what the model sees.
|
||||
type deeplEnvelope struct {
|
||||
Results []deeplTranslation `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// deeplFreeEndpoint is the DeepL free-plan API host. The pro plan uses
|
||||
// api.deepl.com; we default to free because it is the public default
|
||||
// in the Python tool. Override with deeplEndpoint for tests.
|
||||
var deeplFreeEndpoint = "https://api-free.deepl.com/v2/translate"
|
||||
|
||||
// deeplProEndpoint is the DeepL pro-plan API host.
|
||||
var deeplProEndpoint = "https://api.deepl.com/v2/translate"
|
||||
|
||||
// DeepLTool is the Phase 3 batch 3 implementation of the DeepL
|
||||
// translation tool (plan §2.11.4 row 5, §5 Phase 3 第 3 批). It POSTs
|
||||
// a translation request to the DeepL /v2/translate endpoint via the
|
||||
// shared HTTPHelper.
|
||||
type DeepLTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewDeepLTool returns a DeepLTool using the default HTTPHelper.
|
||||
func NewDeepLTool() *DeepLTool {
|
||||
return NewDeepLToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewDeepLToolWith returns a DeepLTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewDeepLToolWith(h *HTTPHelper) *DeepLTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &DeepLTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (d *DeepLTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: deeplToolName,
|
||||
Desc: deeplToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"api_key": {
|
||||
Type: schema.String,
|
||||
Desc: "DeepL API authentication key. Free keys end in ':fx'.",
|
||||
Required: true,
|
||||
},
|
||||
"text": {
|
||||
Type: schema.String,
|
||||
Desc: "Text to translate.",
|
||||
Required: true,
|
||||
},
|
||||
"source_lang": {
|
||||
Type: schema.String,
|
||||
Desc: `Source language code (e.g. "EN", "DE"). Defaults to "EN".`,
|
||||
Required: false,
|
||||
},
|
||||
"target_lang": {
|
||||
Type: schema.String,
|
||||
Desc: `Target language code (e.g. "ZH", "EN-US"). Defaults to "ZH".`,
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildDeepLFormBody composes the application/x-www-form-urlencoded
|
||||
// body that the DeepL /v2/translate endpoint expects. Centralized so
|
||||
// the test suite can verify field encoding.
|
||||
func buildDeepLFormBody(text, sourceLang, targetLang string) string {
|
||||
form := url.Values{}
|
||||
form.Set("text", text)
|
||||
if sourceLang != "" {
|
||||
form.Set("source_lang", strings.ToUpper(sourceLang))
|
||||
}
|
||||
if targetLang != "" {
|
||||
form.Set("target_lang", strings.ToUpper(targetLang))
|
||||
}
|
||||
return form.Encode()
|
||||
}
|
||||
|
||||
// InvokableRun performs the DeepL translation.
|
||||
func (d *DeepLTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p deeplParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return deeplErrJSON(fmt.Errorf("deepl: parse arguments: %w", err)),
|
||||
fmt.Errorf("deepl: parse arguments: %w", err)
|
||||
}
|
||||
if p.APIKey == "" {
|
||||
return deeplErrJSON(fmt.Errorf("api_key is required")),
|
||||
fmt.Errorf("deepl: api_key is required")
|
||||
}
|
||||
if strings.TrimSpace(p.Text) == "" {
|
||||
return deeplErrJSON(fmt.Errorf("text is required")),
|
||||
fmt.Errorf("deepl: text is required")
|
||||
}
|
||||
if p.SourceLang == "" {
|
||||
p.SourceLang = "EN"
|
||||
}
|
||||
if p.TargetLang == "" {
|
||||
p.TargetLang = "ZH"
|
||||
}
|
||||
|
||||
endpoint := deeplFreeEndpoint
|
||||
if !strings.HasSuffix(p.APIKey, ":fx") {
|
||||
// non-:fx keys are pro plan keys; route to the pro endpoint.
|
||||
endpoint = deeplProEndpoint
|
||||
}
|
||||
|
||||
body := buildDeepLFormBody(p.Text, p.SourceLang, p.TargetLang)
|
||||
headers := map[string]string{
|
||||
"Authorization": "DeepL-Auth-Key " + p.APIKey,
|
||||
}
|
||||
|
||||
resp, err := d.helper.Do(ctx, http.MethodPost, endpoint, body,
|
||||
"application/x-www-form-urlencoded", headers)
|
||||
if err != nil {
|
||||
return deeplErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return deeplErrJSON(fmt.Errorf("deepl: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("deepl: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw deeplResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return deeplErrJSON(fmt.Errorf("deepl: decode response: %w", err)),
|
||||
fmt.Errorf("deepl: decode response: %w", err)
|
||||
}
|
||||
return deeplJSON(deeplEnvelope{Results: raw.Translations}), nil
|
||||
}
|
||||
|
||||
func deeplJSON(env deeplEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"deepl: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func deeplErrJSON(err error) string {
|
||||
return deeplJSON(deeplEnvelope{Error: err.Error()})
|
||||
}
|
||||
155
internal/agent/tool/deepl_test.go
Normal file
155
internal/agent/tool/deepl_test.go
Normal file
@@ -0,0 +1,155 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDeepL_BuildRequest(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var gotMethod, gotAuth, gotCT, gotPath string
|
||||
var gotForm url.Values
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotMethod = r.Method
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
gotCT = r.Header.Get("Content-Type")
|
||||
gotPath = r.URL.Path
|
||||
_ = r.ParseForm()
|
||||
gotForm = r.PostForm
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"translations": [
|
||||
{"text":"Hallo Welt","detected_source_language":"EN"}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewDeepLToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"api_key":"key-xyz:fx","text":"Hello world","source_lang":"en","target_lang":"de"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
if gotMethod != http.MethodPost {
|
||||
t.Errorf("method = %q, want POST", gotMethod)
|
||||
}
|
||||
if !strings.HasSuffix(gotPath, "/v2/translate") {
|
||||
t.Errorf("path = %q, want to end with /v2/translate", gotPath)
|
||||
}
|
||||
if gotAuth != "DeepL-Auth-Key key-xyz:fx" {
|
||||
t.Errorf("Authorization = %q, want DeepL-Auth-Key key-xyz:fx", gotAuth)
|
||||
}
|
||||
if !strings.HasPrefix(gotCT, "application/x-www-form-urlencoded") {
|
||||
t.Errorf("Content-Type = %q, want application/x-www-form-urlencoded", gotCT)
|
||||
}
|
||||
if gotForm.Get("text") != "Hello world" {
|
||||
t.Errorf("form.text = %q, want Hello world", gotForm.Get("text"))
|
||||
}
|
||||
// The endpoint routing should send :fx keys to the free endpoint.
|
||||
if gotForm.Get("source_lang") != "EN" {
|
||||
t.Errorf("form.source_lang = %q, want EN (uppercased)", gotForm.Get("source_lang"))
|
||||
}
|
||||
if gotForm.Get("target_lang") != "DE" {
|
||||
t.Errorf("form.target_lang = %q, want DE (uppercased)", gotForm.Get("target_lang"))
|
||||
}
|
||||
|
||||
var env deeplEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 1 {
|
||||
t.Fatalf("Results len = %d, want 1", len(env.Results))
|
||||
}
|
||||
if env.Results[0].Text != "Hallo Welt" {
|
||||
t.Errorf("Results[0].Text = %q, want Hallo Welt", env.Results[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepL_DefaultLanguages(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var gotForm url.Values
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_ = r.ParseForm()
|
||||
gotForm = r.PostForm
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"translations":[{"text":"你好","detected_source_language":"EN"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewDeepLToolWith(helper)
|
||||
if _, err := tool.InvokableRun(context.Background(),
|
||||
`{"api_key":"x:fx","text":"Hello"}`); err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if gotForm.Get("source_lang") != "EN" {
|
||||
t.Errorf("default source_lang = %q, want EN", gotForm.Get("source_lang"))
|
||||
}
|
||||
if gotForm.Get("target_lang") != "ZH" {
|
||||
t.Errorf("default target_lang = %q, want ZH", gotForm.Get("target_lang"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepL_RequiresAPIKeyAndText(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewDeepLTool()
|
||||
if _, err := tool.InvokableRun(context.Background(),
|
||||
`{"api_key":"","text":"Hello"}`); err == nil {
|
||||
t.Error("expected error for missing api_key")
|
||||
}
|
||||
if _, err := tool.InvokableRun(context.Background(),
|
||||
`{"api_key":"x","text":""}`); err == nil {
|
||||
t.Error("expected error for empty text")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepL_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewDeepLTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "deepl" {
|
||||
t.Errorf("Name = %q, want deepl", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "DeepL") {
|
||||
t.Errorf("Desc = %q, want to mention DeepL", info.Desc)
|
||||
}
|
||||
}
|
||||
206
internal/agent/tool/duckduckgo.go
Normal file
206
internal/agent/tool/duckduckgo.go
Normal file
@@ -0,0 +1,206 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const duckduckgoToolName = "duckduckgo"
|
||||
|
||||
const duckduckgoToolDescription = "Search DuckDuckGo's Instant Answer API. Returns the abstract text and a list of related topics."
|
||||
|
||||
// duckduckgoParams is the JSON shape the model sends into InvokableRun.
|
||||
// max_results caps the number of RelatedTopics returned; the upstream
|
||||
// endpoint itself is uncapped and uses heuristic ranking.
|
||||
type duckduckgoParams struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
// duckduckgoTopic is one element of the upstream `RelatedTopics` array.
|
||||
// DuckDuckGo returns a recursive tree: a topic may itself have
|
||||
// `Topics`, and may be a "name"/"value" pair (no URL) or a regular
|
||||
// hit with `FirstURL`. We flatten one level of nesting.
|
||||
type duckduckgoTopic struct {
|
||||
Text string `json:"text,omitempty"`
|
||||
FirstURL string `json:"first_url,omitempty"`
|
||||
Topics []duckduckgoTopic `json:"topics,omitempty"`
|
||||
}
|
||||
|
||||
// duckduckgoResponse is the upstream Instant Answer envelope. We only
|
||||
// model the fields we care about.
|
||||
type duckduckgoResponse struct {
|
||||
AbstractText string `json:"abstract_text"`
|
||||
AbstractURL string `json:"abstract_url"`
|
||||
Abstract string `json:"abstract"`
|
||||
RelatedTopics []duckduckgoTopic `json:"related_topics"`
|
||||
}
|
||||
|
||||
// duckduckgoTopicOut is the model-facing topic shape.
|
||||
type duckduckgoTopicOut struct {
|
||||
Text string `json:"text,omitempty"`
|
||||
FirstURL string `json:"first_url,omitempty"`
|
||||
}
|
||||
|
||||
// duckduckgoEnvelope is the JSON shape the model sees.
|
||||
type duckduckgoEnvelope struct {
|
||||
AbstractText string `json:"abstract_text,omitempty"`
|
||||
AbstractURL string `json:"abstract_url,omitempty"`
|
||||
RelatedTopics []duckduckgoTopicOut `json:"related_topics,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// DuckDuckGoTool is the Phase 3 batch 2 implementation of the DuckDuckGo
|
||||
// Instant Answer tool (plan §2.11.4 row 6, §5 Phase 3 第 2 批). It
|
||||
// performs a GET against the public Instant Answer endpoint using the
|
||||
// shared HTTPHelper and returns the abstract + a flat list of related
|
||||
// topics.
|
||||
type DuckDuckGoTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewDuckDuckGoTool returns a DuckDuckGoTool using the default HTTPHelper.
|
||||
func NewDuckDuckGoTool() *DuckDuckGoTool {
|
||||
return NewDuckDuckGoToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewDuckDuckGoToolWith returns a DuckDuckGoTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewDuckDuckGoToolWith(h *HTTPHelper) *DuckDuckGoTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &DuckDuckGoTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (d *DuckDuckGoTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: duckduckgoToolName,
|
||||
Desc: duckduckgoToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query",
|
||||
Required: true,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of related topics to return. Defaults to 5.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildDuckDuckGoURL constructs the Instant Answer API URL.
|
||||
func buildDuckDuckGoURL(query string) string {
|
||||
q := url.Values{}
|
||||
q.Set("q", query)
|
||||
q.Set("format", "json")
|
||||
q.Set("no_html", "1")
|
||||
q.Set("skip_disambig", "1")
|
||||
return "https://api.duckduckgo.com/?" + q.Encode()
|
||||
}
|
||||
|
||||
// flattenDuckDuckGoTopics flattens the upstream topic tree (which can
|
||||
// have arbitrary nesting) into a list, dropping entries without a URL.
|
||||
func flattenDuckDuckGoTopics(in []duckduckgoTopic) []duckduckgoTopicOut {
|
||||
var out []duckduckgoTopicOut
|
||||
for _, t := range in {
|
||||
if t.FirstURL != "" && t.Text != "" {
|
||||
out = append(out, duckduckgoTopicOut{Text: t.Text, FirstURL: t.FirstURL})
|
||||
}
|
||||
// recurse one level; upstream nests categories here
|
||||
if len(t.Topics) > 0 {
|
||||
out = append(out, flattenDuckDuckGoTopics(t.Topics)...)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// InvokableRun performs the DuckDuckGo Instant Answer query.
|
||||
func (d *DuckDuckGoTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p duckduckgoParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return duckduckgoErrJSON(fmt.Errorf("duckduckgo: parse arguments: %w", err)),
|
||||
fmt.Errorf("duckduckgo: parse arguments: %w", err)
|
||||
}
|
||||
if p.Query == "" {
|
||||
return duckduckgoErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("duckduckgo: query is required")
|
||||
}
|
||||
if p.MaxResults <= 0 {
|
||||
p.MaxResults = 5
|
||||
}
|
||||
|
||||
endpoint := buildDuckDuckGoURL(p.Query)
|
||||
resp, err := d.helper.Do(ctx, http.MethodGet, endpoint, "", "", nil)
|
||||
if err != nil {
|
||||
return duckduckgoErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return duckduckgoErrJSON(fmt.Errorf("duckduckgo: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("duckduckgo: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw duckduckgoResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return duckduckgoErrJSON(fmt.Errorf("duckduckgo: decode response: %w", err)),
|
||||
fmt.Errorf("duckduckgo: decode response: %w", err)
|
||||
}
|
||||
|
||||
// Prefer AbstractText (rich) over Abstract (plain), as the Python
|
||||
// tool did historically.
|
||||
abstract := raw.AbstractText
|
||||
if abstract == "" {
|
||||
abstract = raw.Abstract
|
||||
}
|
||||
|
||||
topics := flattenDuckDuckGoTopics(raw.RelatedTopics)
|
||||
if len(topics) > p.MaxResults {
|
||||
topics = topics[:p.MaxResults]
|
||||
}
|
||||
|
||||
return duckduckgoJSON(duckduckgoEnvelope{
|
||||
AbstractText: abstract,
|
||||
AbstractURL: raw.AbstractURL,
|
||||
RelatedTopics: topics,
|
||||
}), nil
|
||||
}
|
||||
|
||||
func duckduckgoJSON(env duckduckgoEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"duckduckgo: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func duckduckgoErrJSON(err error) string {
|
||||
return duckduckgoJSON(duckduckgoEnvelope{Error: err.Error()})
|
||||
}
|
||||
256
internal/agent/tool/duckduckgo_test.go
Normal file
256
internal/agent/tool/duckduckgo_test.go
Normal file
@@ -0,0 +1,256 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/flow/agent/react"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
func TestDuckDuckGo_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got := buildDuckDuckGoURL("rag flow")
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != "api.duckduckgo.com" {
|
||||
t.Errorf("host = %q, want api.duckduckgo.com", u.Host)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("q") != "rag flow" {
|
||||
t.Errorf("q = %q, want 'rag flow' (no pre-encoding)", q.Get("q"))
|
||||
}
|
||||
if q.Get("format") != "json" {
|
||||
t.Errorf("format = %q, want json", q.Get("format"))
|
||||
}
|
||||
if q.Get("no_html") != "1" {
|
||||
t.Errorf("no_html = %q, want 1", q.Get("no_html"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuckDuckGo_ParseTopics(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
// upstream returns a recursive tree; include a category
|
||||
// (`Topics` non-empty, no `FirstURL`) and a leaf hit.
|
||||
_, _ = w.Write([]byte(`{
|
||||
"abstract_text": "RAGFlow is an open-source RAG engine.",
|
||||
"abstract_url": "https://ragflow.io",
|
||||
"related_topics": [
|
||||
{
|
||||
"text": "Category: Technology",
|
||||
"topics": [
|
||||
{"text": "Open source", "first_url": "https://example.com/os"},
|
||||
{"text": "Search engines", "first_url": "https://example.com/se"}
|
||||
]
|
||||
},
|
||||
{"text": "GitHub", "first_url": "https://github.com/infiniflow/ragflow"}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewDuckDuckGoToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(), `{"query":"ragflow","max_results":10}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env duckduckgoEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if env.AbstractText != "RAGFlow is an open-source RAG engine." {
|
||||
t.Errorf("AbstractText = %q, want the upstream abstract", env.AbstractText)
|
||||
}
|
||||
if env.AbstractURL != "https://ragflow.io" {
|
||||
t.Errorf("AbstractURL = %q, want https://ragflow.io", env.AbstractURL)
|
||||
}
|
||||
if len(env.RelatedTopics) != 3 {
|
||||
t.Fatalf("RelatedTopics len = %d, want 3 (category child leaves + direct hit)", len(env.RelatedTopics))
|
||||
}
|
||||
wantURLs := map[string]bool{
|
||||
"https://example.com/os": false,
|
||||
"https://example.com/se": false,
|
||||
"https://github.com/infiniflow/ragflow": false,
|
||||
}
|
||||
for _, t2 := range env.RelatedTopics {
|
||||
if _, ok := wantURLs[t2.FirstURL]; ok {
|
||||
wantURLs[t2.FirstURL] = true
|
||||
}
|
||||
}
|
||||
for u, seen := range wantURLs {
|
||||
if !seen {
|
||||
t.Errorf("missing topic URL %q in flattened result", u)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuckDuckGo_RespectsMaxResults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"related_topics": [
|
||||
{"text":"A","first_url":"https://a"},
|
||||
{"text":"B","first_url":"https://b"},
|
||||
{"text":"C","first_url":"https://c"},
|
||||
{"text":"D","first_url":"https://d"}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewDuckDuckGoToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(), `{"query":"x","max_results":2}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env duckduckgoEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v", jerr)
|
||||
}
|
||||
if len(env.RelatedTopics) != 2 {
|
||||
t.Errorf("RelatedTopics len = %d, want 2 (capped by max_results)", len(env.RelatedTopics))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuckDuckGo_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewDuckDuckGoTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "duckduckgo" {
|
||||
t.Errorf("Name = %q, want duckduckgo", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "DuckDuckGo") {
|
||||
t.Errorf("Desc = %q, want to mention DuckDuckGo", info.Desc)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDuckDuckGo_RealReactAgent_ExecutesTool drives a real eino
|
||||
// react.NewAgent with the real DuckDuckGoTool (httptest-stubbed
|
||||
// upstream) and a scripted chat model. Proves the HTTP-based tool is
|
||||
// actually invoked by eino, its JSON envelope is fed back as a
|
||||
// ToolMessage, and the model can ground a final answer in the
|
||||
// retrieved abstract. This is the HTTP-side counterpart to
|
||||
// TestExeSQL_RealReactAgent_ExecutesTool — together they cover the
|
||||
// two distinct wiring patterns (DB-backed vs HTTP-backed) the agent
|
||||
// needs to handle.
|
||||
func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var hitCount int
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hitCount++
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"abstract_text": "RAGFlow is an open-source RAG engine.",
|
||||
"abstract_url": "https://ragflow.io",
|
||||
"related_topics": [
|
||||
{"text": "GitHub", "first_url": "https://github.com/infiniflow/ragflow"}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
realTool := NewDuckDuckGoToolWith(helper)
|
||||
|
||||
mdl := newReactScriptedModel(
|
||||
"duckduckgo",
|
||||
`{"query":"ragflow"}`,
|
||||
"RAGFlow is an open-source RAG engine.",
|
||||
)
|
||||
|
||||
agent, err := react.NewAgent(context.Background(), &react.AgentConfig{
|
||||
ToolCallingModel: mdl,
|
||||
ToolsConfig: compose.ToolsNodeConfig{
|
||||
Tools: []einotool.BaseTool{realTool},
|
||||
},
|
||||
MaxStep: 5,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("react.NewAgent: %v", err)
|
||||
}
|
||||
|
||||
out, err := agent.Generate(context.Background(), []*schema.Message{
|
||||
schema.UserMessage("What is RAGFlow?"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("agent.Generate: %v", err)
|
||||
}
|
||||
if got, want := out.Content, "RAGFlow is an open-source RAG engine."; got != want {
|
||||
t.Errorf("Content = %q, want %q", got, want)
|
||||
}
|
||||
if mdl.turn != 2 {
|
||||
t.Errorf("Generate called %d times, want 2 (tool_call + final)", mdl.turn)
|
||||
}
|
||||
if len(mdl.boundTools) != 1 || mdl.boundTools[0].Name != "duckduckgo" {
|
||||
names := make([]string, 0, len(mdl.boundTools))
|
||||
for _, ti := range mdl.boundTools {
|
||||
names = append(names, ti.Name)
|
||||
}
|
||||
t.Errorf("tools bound to model = %v, want [duckduckgo]", names)
|
||||
}
|
||||
if len(mdl.rounds) < 2 {
|
||||
t.Fatalf("only %d rounds captured, want >= 2", len(mdl.rounds))
|
||||
}
|
||||
var sawToolResult bool
|
||||
for _, msg := range mdl.rounds[1] {
|
||||
if msg.Role == schema.Tool && strings.Contains(msg.Content, "RAGFlow is an open-source RAG engine") {
|
||||
sawToolResult = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !sawToolResult {
|
||||
t.Errorf("round 2 input did not contain a ToolMessage carrying the upstream abstract")
|
||||
}
|
||||
if hitCount == 0 {
|
||||
t.Error("test server was never hit; the tool did not actually call the upstream")
|
||||
}
|
||||
}
|
||||
195
internal/agent/tool/email.go
Normal file
195
internal/agent/tool/email.go
Normal file
@@ -0,0 +1,195 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/smtp"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const emailToolName = "email"
|
||||
|
||||
const emailToolDescription = "Send an email via SMTP. Returns success/failure status."
|
||||
|
||||
// emailParams is the JSON shape the model sends into InvokableRun.
|
||||
type emailParams struct {
|
||||
SMTPHost string `json:"smtp_host"`
|
||||
SMTPPort int `json:"smtp_port"`
|
||||
Username string `json:"username"`
|
||||
Password string `json:"password"`
|
||||
FromAddr string `json:"from_addr"`
|
||||
ToAddrs []string `json:"to_addrs"`
|
||||
Subject string `json:"subject"`
|
||||
Body string `json:"body"`
|
||||
}
|
||||
|
||||
// emailEnvelope is what the model sees.
|
||||
type emailEnvelope struct {
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// EmailTool is the Phase 3 batch 3 implementation of the SMTP email
|
||||
// sender tool (plan §2.11.4 row 7, §5 Phase 3 第 3 批). It composes an
|
||||
// RFC 822 message and submits it via the stdlib net/smtp client. All
|
||||
// authentication modes supported by net/smtp.Auth are available
|
||||
// (PLAIN, LOGIN, CRAM-MD5) by selecting the appropriate creds.
|
||||
type EmailTool struct{}
|
||||
|
||||
// NewEmailTool returns an EmailTool. There is no shared HTTPHelper
|
||||
// (SMTP is not HTTP), so the constructor is the simplest possible.
|
||||
func NewEmailTool() *EmailTool {
|
||||
return &EmailTool{}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (e *EmailTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: emailToolName,
|
||||
Desc: emailToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"smtp_host": {
|
||||
Type: schema.String,
|
||||
Desc: "SMTP server hostname (e.g. smtp.gmail.com).",
|
||||
Required: true,
|
||||
},
|
||||
"smtp_port": {
|
||||
Type: schema.Integer,
|
||||
Desc: "SMTP server port (e.g. 587 for STARTTLS, 465 for implicit TLS).",
|
||||
Required: true,
|
||||
},
|
||||
"username": {
|
||||
Type: schema.String,
|
||||
Desc: "SMTP authentication username. Empty for unauthenticated relay.",
|
||||
Required: false,
|
||||
},
|
||||
"password": {
|
||||
Type: schema.String,
|
||||
Desc: "SMTP authentication password (or app password for Gmail/Yahoo).",
|
||||
Required: false,
|
||||
},
|
||||
"from_addr": {
|
||||
Type: schema.String,
|
||||
Desc: "Sender email address (RFC 5322).",
|
||||
Required: true,
|
||||
},
|
||||
"to_addrs": {
|
||||
Type: schema.Array,
|
||||
Desc: "Recipient email addresses.",
|
||||
Required: true,
|
||||
},
|
||||
"subject": {
|
||||
Type: schema.String,
|
||||
Desc: "Email subject line.",
|
||||
Required: true,
|
||||
},
|
||||
"body": {
|
||||
Type: schema.String,
|
||||
Desc: "Email body (plain text).",
|
||||
Required: true,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildEmailMessage composes the RFC 822 wire format: headers + blank
|
||||
// line + body. Extracted so tests can verify subject / recipient
|
||||
// inclusion without opening a real socket.
|
||||
func buildEmailMessage(from string, to []string, subject, body string) []byte {
|
||||
var b strings.Builder
|
||||
b.WriteString("From: " + from + "\r\n")
|
||||
b.WriteString("To: " + strings.Join(to, ", ") + "\r\n")
|
||||
b.WriteString("Subject: " + subject + "\r\n")
|
||||
b.WriteString("MIME-Version: 1.0\r\n")
|
||||
b.WriteString("Content-Type: text/plain; charset=UTF-8\r\n")
|
||||
b.WriteString("\r\n")
|
||||
b.WriteString(body)
|
||||
b.WriteString("\r\n")
|
||||
return []byte(b.String())
|
||||
}
|
||||
|
||||
// InvokableRun sends the email. We delegate to smtp.SendMail which
|
||||
// handles EHLO, STARTTLS, and AUTH transparently when an *smtp.Auth is
|
||||
// supplied; with nil auth it sends unauthenticated.
|
||||
func (e *EmailTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p emailParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return emailErrJSON(fmt.Errorf("email: parse arguments: %w", err)),
|
||||
fmt.Errorf("email: parse arguments: %w", err)
|
||||
}
|
||||
if err := validateEmailParams(&p); err != nil {
|
||||
return emailErrJSON(err), err
|
||||
}
|
||||
|
||||
addr := fmt.Sprintf("%s:%d", p.SMTPHost, p.SMTPPort)
|
||||
msg := buildEmailMessage(p.FromAddr, p.ToAddrs, p.Subject, p.Body)
|
||||
|
||||
var auth smtp.Auth
|
||||
if p.Username != "" {
|
||||
auth = smtp.PlainAuth("", p.Username, p.Password, p.SMTPHost)
|
||||
}
|
||||
|
||||
if err := smtp.SendMail(addr, auth, p.FromAddr, p.ToAddrs, msg); err != nil {
|
||||
return emailErrJSON(fmt.Errorf("email: send: %w", err)),
|
||||
fmt.Errorf("email: send: %w", err)
|
||||
}
|
||||
|
||||
// Honor context cancellation if the caller passed a deadline. The
|
||||
// underlying smtp.SendMail is blocking, so we check after the call;
|
||||
// a stricter impl would select on ctx.Done() around the call.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return emailErrJSON(err), err
|
||||
}
|
||||
return emailJSON(emailEnvelope{OK: true}), nil
|
||||
}
|
||||
|
||||
// validateEmailParams guards against obviously broken inputs. The
|
||||
// upstream SMTP server will give a better error for malformed
|
||||
// addresses, but the common case (empty / missing) is caught here.
|
||||
func validateEmailParams(p *emailParams) error {
|
||||
switch {
|
||||
case p.SMTPHost == "":
|
||||
return fmt.Errorf("email: smtp_host is required")
|
||||
case p.SMTPPort <= 0 || p.SMTPPort > 65535:
|
||||
return fmt.Errorf("email: smtp_port must be in [1, 65535]")
|
||||
case p.FromAddr == "":
|
||||
return fmt.Errorf("email: from_addr is required")
|
||||
case len(p.ToAddrs) == 0:
|
||||
return fmt.Errorf("email: to_addrs is required and must be non-empty")
|
||||
case p.Subject == "":
|
||||
return fmt.Errorf("email: subject is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func emailJSON(env emailEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"email: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func emailErrJSON(err error) string {
|
||||
return emailJSON(emailEnvelope{Error: err.Error()})
|
||||
}
|
||||
229
internal/agent/tool/email_test.go
Normal file
229
internal/agent/tool/email_test.go
Normal file
@@ -0,0 +1,229 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestEmail_BuildMessage(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
msg := buildEmailMessage(
|
||||
"alice@example.com",
|
||||
[]string{"bob@example.com", "carol@example.com"},
|
||||
"Hello, world",
|
||||
"Body of the message.",
|
||||
)
|
||||
|
||||
s := string(msg)
|
||||
for _, want := range []string{
|
||||
"From: alice@example.com",
|
||||
"To: bob@example.com, carol@example.com",
|
||||
"Subject: Hello, world",
|
||||
"Content-Type: text/plain; charset=UTF-8",
|
||||
"Body of the message.",
|
||||
} {
|
||||
if !strings.Contains(s, want) {
|
||||
t.Errorf("message missing %q\n--- message ---\n%s\n---", want, s)
|
||||
}
|
||||
}
|
||||
// RFC 822 mandates a blank line between headers and body.
|
||||
if !strings.Contains(s, "\r\n\r\n") {
|
||||
t.Errorf("message missing blank line between headers and body\n%s", s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmail_SendAgainstMockSMTP(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Spin up a minimal SMTP server: read commands, respond 250 to
|
||||
// everything, and copy the DATA payload bytes so the test can
|
||||
// inspect them.
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("Listen: %v", err)
|
||||
}
|
||||
defer ln.Close()
|
||||
|
||||
var receivedData strings.Builder
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
|
||||
reader := bufio.NewReader(conn)
|
||||
writer := bufio.NewWriter(conn)
|
||||
// Greeting
|
||||
_, _ = writer.WriteString("220 mock-smtp ready\r\n")
|
||||
_ = writer.Flush()
|
||||
inData := false
|
||||
for {
|
||||
line, err := reader.ReadString('\n')
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
up := strings.ToUpper(strings.TrimSpace(line))
|
||||
switch {
|
||||
case strings.HasPrefix(up, "EHLO"), strings.HasPrefix(up, "HELO"):
|
||||
_, _ = writer.WriteString("250-mock-smtp\r\n250 OK\r\n")
|
||||
_ = writer.Flush()
|
||||
case strings.HasPrefix(up, "MAIL FROM:"), strings.HasPrefix(up, "RCPT TO:"):
|
||||
_, _ = writer.WriteString("250 OK\r\n")
|
||||
_ = writer.Flush()
|
||||
case strings.HasPrefix(up, "DATA"):
|
||||
_, _ = writer.WriteString("354 End data with <CR><LF>.<CR><LF>\r\n")
|
||||
_ = writer.Flush()
|
||||
inData = true
|
||||
case inData && strings.TrimSpace(line) == ".":
|
||||
_, _ = writer.WriteString("250 Queued\r\n")
|
||||
_ = writer.Flush()
|
||||
inData = false
|
||||
case inData:
|
||||
receivedData.WriteString(line)
|
||||
case strings.HasPrefix(up, "QUIT"):
|
||||
_, _ = writer.WriteString("221 Bye\r\n")
|
||||
_ = writer.Flush()
|
||||
return
|
||||
default:
|
||||
_, _ = writer.WriteString("250 OK\r\n")
|
||||
_ = writer.Flush()
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
host, port, err := net.SplitHostPort(ln.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatalf("SplitHostPort: %v", err)
|
||||
}
|
||||
_ = host
|
||||
var portInt int
|
||||
_, _ = fmt.Sscanf(port, "%d", &portInt)
|
||||
|
||||
tool := NewEmailTool()
|
||||
args := map[string]any{
|
||||
"smtp_host": "127.0.0.1",
|
||||
"smtp_port": portInt,
|
||||
"from_addr": "alice@example.com",
|
||||
"to_addrs": []string{"bob@example.com"},
|
||||
"subject": "Test Subject",
|
||||
"body": "Test body content.",
|
||||
}
|
||||
argsJSON, _ := json.Marshal(args)
|
||||
out, err := tool.InvokableRun(context.Background(), string(argsJSON))
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env emailEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if !env.OK {
|
||||
t.Errorf("OK = false, want true")
|
||||
}
|
||||
|
||||
// Wait for the mock server to finish.
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("mock SMTP server did not close in time")
|
||||
}
|
||||
|
||||
if !strings.Contains(receivedData.String(), "Subject: Test Subject") {
|
||||
t.Errorf("mock server did not receive subject\n--- data ---\n%s\n---",
|
||||
receivedData.String())
|
||||
}
|
||||
if !strings.Contains(receivedData.String(), "bob@example.com") {
|
||||
t.Errorf("mock server did not receive recipient\n--- data ---\n%s\n---",
|
||||
receivedData.String())
|
||||
}
|
||||
if !strings.Contains(receivedData.String(), "Test body content.") {
|
||||
t.Errorf("mock server did not receive body\n--- data ---\n%s\n---",
|
||||
receivedData.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmail_RequiresFields(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewEmailTool()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
args string
|
||||
wantErr string
|
||||
}{
|
||||
{
|
||||
name: "missing smtp_host",
|
||||
args: `{"smtp_port":587,"from_addr":"a@b","to_addrs":["c@d"],"subject":"s","body":"b"}`,
|
||||
wantErr: "smtp_host",
|
||||
},
|
||||
{
|
||||
name: "missing to_addrs",
|
||||
args: `{"smtp_host":"x","smtp_port":587,"from_addr":"a@b","subject":"s","body":"b"}`,
|
||||
wantErr: "to_addrs",
|
||||
},
|
||||
{
|
||||
name: "bad smtp_port",
|
||||
args: `{"smtp_host":"x","smtp_port":0,"from_addr":"a@b","to_addrs":["c@d"],"subject":"s","body":"b"}`,
|
||||
wantErr: "smtp_port",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
_, err := tool.InvokableRun(context.Background(), tc.args)
|
||||
if err == nil {
|
||||
t.Fatalf("expected error for %s", tc.name)
|
||||
}
|
||||
if !strings.Contains(err.Error(), tc.wantErr) {
|
||||
t.Errorf("err = %v, want to contain %q", err, tc.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmail_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewEmailTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "email" {
|
||||
t.Errorf("Name = %q, want email", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "SMTP") {
|
||||
t.Errorf("Desc = %q, want to mention SMTP", info.Desc)
|
||||
}
|
||||
}
|
||||
574
internal/agent/tool/exesql.go
Normal file
574
internal/agent/tool/exesql.go
Normal file
@@ -0,0 +1,574 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package tool — ExeSQL tool (Phase 3 batch 1).
|
||||
//
|
||||
// ExeSQL lets an Agent component execute a SQL statement on a
|
||||
// user-configured external database and return the rows as
|
||||
// column→value maps. It mirrors the Python `agent/tools/exesql.py`
|
||||
// semantics: per-invocation open/close of a fresh `database/sql` connection
|
||||
// scoped to the tool's params (host/port/database/username/password),
|
||||
// no ORM, no GORM — those layers are for RAGFlow's own metadata DB
|
||||
// (`internal/dao`) and would be the wrong abstraction here.
|
||||
//
|
||||
// Why `database/sql` (not GORM, not sqlx):
|
||||
// - The Python equivalent uses `pymysql.connect()` / `psycopg2.connect()`
|
||||
// directly with no ORM in the path.
|
||||
// - ExeSQL returns `[]map[string]any` (dynamic schema, LLM-supplied
|
||||
// SQL), so there is nothing to struct-scan — sqlx's headline
|
||||
// feature would be unused.
|
||||
// - GORM is for object-relational mapping of RAGFlow's metadata
|
||||
// tables. Reusing `internal/dao`'s GORM here would couple the
|
||||
// tool to RAGFlow's own DB pool and require an `internal/dao`
|
||||
// connection for a tool that has no business touching RAGFlow's
|
||||
// metadata at all.
|
||||
// - `database/sql` is stdlib, no extra runtime dependency beyond
|
||||
// the per-driver package (`go-sql-driver/mysql`, `lib/pq`,
|
||||
// `denisenkom/go-mssqldb`).
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
// SQL drivers — registered via their init() side effects.
|
||||
_ "github.com/denisenkom/go-mssqldb"
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
_ "github.com/lib/pq"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// ExeSQL-specific errors. The previous Phase 3 batch 1 stub returned
|
||||
// ErrExeSQLDAOMissing because the implementation was a no-op pending
|
||||
// the (now rejected) GORM wiring. With the real `database/sql`
|
||||
// implementation in place, the error surface is:
|
||||
// - ErrExeSQLNotSelect: SQL failed the SELECT-only safety filter.
|
||||
// - ErrExeSQLNoCredentials: the tool has no db_type/host/etc. set
|
||||
// (caller forgot to wire the connection params).
|
||||
// - ErrExeSQLUnsupportedDB: db_type is one we have not ported yet
|
||||
// (trino, IBM DB2).
|
||||
var (
|
||||
ErrExeSQLNotSelect = errors.New(
|
||||
"ExeSQL: only SELECT statements are allowed; " +
|
||||
"INSERT/UPDATE/DELETE/DDL is rejected for safety",
|
||||
)
|
||||
ErrExeSQLNoCredentials = errors.New(
|
||||
"ExeSQL: connection params not configured (db_type/host/port/database/username/password)",
|
||||
)
|
||||
ErrExeSQLUnsupportedDB = errors.New(
|
||||
"ExeSQL: db_type not yet supported in Go port (trino, IBM DB2 are pending)",
|
||||
)
|
||||
)
|
||||
|
||||
const (
|
||||
exesqlToolName = "execute_sql"
|
||||
exesqlToolDescription = "This is a tool that can execute SQL."
|
||||
|
||||
exesqlDefaultMaxRecords = 1024
|
||||
exesqlDefaultTimeout = 60 * time.Second
|
||||
)
|
||||
|
||||
// exesqlConnParams captures the user-supplied DB connection details.
|
||||
// These are tool-level config (set on the canvas node, not exposed
|
||||
// to the LLM at function-call time), matching the Python ExeSQLParam
|
||||
// fields. The LLM only sees `sql` and optional `database` in args.
|
||||
type exesqlConnParams struct {
|
||||
DBType string // mysql | postgres | mariadb | mssql | oceanbase
|
||||
Database string
|
||||
Username string
|
||||
Host string
|
||||
Port int
|
||||
Password string
|
||||
MaxRecords int
|
||||
}
|
||||
|
||||
// exesqlArgs is the JSON shape the model sends in. Matches the Python
|
||||
// ExeSQLParam ToolMeta (sql is required, database is optional).
|
||||
type exesqlArgs struct {
|
||||
SQL string `json:"sql"`
|
||||
Database string `json:"database,omitempty"`
|
||||
}
|
||||
|
||||
// exesqlResult is the JSON envelope returned to the model. The shape
|
||||
// matches the Python tool's `rows` / `_ERROR` output convention so
|
||||
// downstream nodes can pattern-match unchanged. `Columns` lets the
|
||||
// caller preserve order even when `Rows` is empty.
|
||||
type exesqlResult struct {
|
||||
Columns []string `json:"columns,omitempty"`
|
||||
Rows []map[string]any `json:"rows,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// exesqlDialer abstracts `sql.Open` for test injection. Production code
|
||||
// uses defaultExeSQLDialer (== sql.Open). Tests inject a dialer that
|
||||
// returns a `*sql.DB` backed by DATA-DOG/go-sqlmock so no real DB is
|
||||
// required.
|
||||
//
|
||||
// Returning a *sql.DB (not a Tx) means each statement gets its own
|
||||
// fresh conn from the pool, matching the Python tool's
|
||||
// open → execute → close lifecycle.
|
||||
type exesqlDialer func(driver, dsn string) (*sql.DB, error)
|
||||
|
||||
func defaultExeSQLDialer(driver, dsn string) (*sql.DB, error) {
|
||||
db, err := sql.Open(driver, dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// The Python tool does not pool — every call is a fresh connect.
|
||||
// Match that: set MaxOpenConns=1 and Close() in InvokableRun so
|
||||
// we never leak a connection. Timeouts are honoured by the caller's
|
||||
// context.
|
||||
db.SetMaxOpenConns(1)
|
||||
return db, nil
|
||||
}
|
||||
|
||||
// ExeSQLTool is the Phase 3 batch 1 implementation of the ExeSQL tool.
|
||||
// It validates SELECT-only at the parser level and executes the
|
||||
// statement against a user-configured external DB via `database/sql`.
|
||||
type ExeSQLTool struct {
|
||||
conn exesqlConnParams
|
||||
dialer exesqlDialer
|
||||
}
|
||||
|
||||
// NewExeSQLTool returns an ExeSQLTool wired to the given connection
|
||||
// params. The dialer defaults to `sql.Open`; tests can pass a
|
||||
// sqlmock-backed dialer via WithExeSQLDialer.
|
||||
func NewExeSQLTool(conn exesqlConnParams) *ExeSQLTool {
|
||||
if conn.MaxRecords <= 0 {
|
||||
conn.MaxRecords = exesqlDefaultMaxRecords
|
||||
}
|
||||
return &ExeSQLTool{
|
||||
conn: conn,
|
||||
dialer: defaultExeSQLDialer,
|
||||
}
|
||||
}
|
||||
|
||||
// WithExeSQLDialer swaps the connection opener. Returns the tool for
|
||||
// chaining. Used by tests; production code should leave it alone.
|
||||
func (e *ExeSQLTool) WithExeSQLDialer(d exesqlDialer) *ExeSQLTool {
|
||||
if d != nil {
|
||||
e.dialer = d
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model. Mirrors the
|
||||
// Python ExeSQLParam ToolMeta: only `sql` (and optional `database`)
|
||||
// are visible to the LLM. Connection params are not exposed here —
|
||||
// they're set on the tool instance, matching the Python convention
|
||||
// where ExeSQLParam fields like `db_type` / `host` are tool
|
||||
// configuration, not function-call arguments.
|
||||
func (e *ExeSQLTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: exesqlToolName,
|
||||
Desc: exesqlToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"sql": {
|
||||
Type: schema.String,
|
||||
Desc: "The SQL statement to execute. Must be a SELECT (read-only).",
|
||||
Required: true,
|
||||
},
|
||||
"database": {
|
||||
Type: schema.String,
|
||||
Desc: "Optional target database / schema name. Overrides the tool's configured DB.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// InvokableRun validates the SQL, opens a fresh connection scoped to
|
||||
// the tool's params, executes each semicolon-separated statement, and
|
||||
// returns the rows. Per-statement errors do not abort the node: they
|
||||
// are accumulated in the `Errors` slice of the response (the Python
|
||||
// tool does the same — `sql_res.append({"content": msg})`).
|
||||
func (e *ExeSQLTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
|
||||
if argumentsInJSON == "" {
|
||||
return exesqlErrorResult(errors.New("exesql: empty arguments")), errors.New("exesql: empty arguments")
|
||||
}
|
||||
var args exesqlArgs
|
||||
if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil {
|
||||
return exesqlErrorResult(fmt.Errorf("exesql: parse arguments: %w", err)),
|
||||
fmt.Errorf("exesql: parse arguments: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(args.SQL) == "" {
|
||||
return exesqlErrorResult(errors.New("exesql: empty sql")), errors.New("exesql: empty sql")
|
||||
}
|
||||
if !isSelectStatement(args.SQL) {
|
||||
return exesqlErrorResult(ErrExeSQLNotSelect), ErrExeSQLNotSelect
|
||||
}
|
||||
|
||||
// Honor the per-call `database` override if the model supplied one;
|
||||
// fall back to the tool's configured DB.
|
||||
conn := e.conn
|
||||
if args.Database != "" {
|
||||
conn.Database = args.Database
|
||||
}
|
||||
if err := conn.check(); err != nil {
|
||||
return exesqlErrorResult(err), err
|
||||
}
|
||||
|
||||
driver, dsn, err := exesqlDriverAndDSN(conn)
|
||||
if err != nil {
|
||||
return exesqlErrorResult(err), err
|
||||
}
|
||||
|
||||
db, err := e.dialer(driver, dsn)
|
||||
if err != nil {
|
||||
return exesqlErrorResult(fmt.Errorf("exesql: open %s: %w", driver, err)),
|
||||
fmt.Errorf("exesql: open %s: %w", driver, err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
// Apply a wall-clock timeout if the caller did not provide one.
|
||||
if _, ok := ctx.Deadline(); !ok {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, exesqlDefaultTimeout)
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
if err := db.PingContext(ctx); err != nil {
|
||||
return exesqlErrorResult(fmt.Errorf("exesql: ping: %w", err)),
|
||||
fmt.Errorf("exesql: ping: %w", err)
|
||||
}
|
||||
|
||||
res, err := exesqlExecute(ctx, db, args.SQL, conn.MaxRecords)
|
||||
if err != nil {
|
||||
return exesqlErrorResult(err), err
|
||||
}
|
||||
return exesqlMarshalResult(res)
|
||||
}
|
||||
|
||||
// exesqlExecute splits the SQL on `;` (Python parity) and runs each
|
||||
// statement independently. A failing statement is recorded as an
|
||||
// error entry but does not abort subsequent statements — this is
|
||||
// the same isolation guarantee the Python tool provides so that
|
||||
// earlier results survive a bad statement later in the batch.
|
||||
func exesqlExecute(ctx context.Context, db *sql.DB, sqlText string, maxRows int) (*exesqlResult, error) {
|
||||
stmts := splitSQLStatements(sqlText)
|
||||
res := &exesqlResult{}
|
||||
for _, stmt := range stmts {
|
||||
stmt = stripChunkIDMarkers(stmt)
|
||||
stmt = strings.TrimSpace(stmt)
|
||||
if stmt == "" {
|
||||
continue
|
||||
}
|
||||
cols, rows, err := exesqlQueryOne(ctx, db, stmt, maxRows)
|
||||
if err != nil {
|
||||
// Keep going; the Python tool appends {"content": msg}
|
||||
// and continues with the next statement.
|
||||
res.Rows = append(res.Rows, map[string]any{
|
||||
"content": "SQL Execution Failed: " + stmt + "\n" + err.Error(),
|
||||
})
|
||||
continue
|
||||
}
|
||||
if len(res.Columns) == 0 && len(cols) > 0 {
|
||||
res.Columns = cols
|
||||
}
|
||||
res.Rows = append(res.Rows, rows...)
|
||||
}
|
||||
if len(res.Rows) == 0 {
|
||||
// Mirror the Python tool's "no record" sentinel so downstream
|
||||
// nodes (VariableAggregator, Message) can match on it.
|
||||
// Trigger on zero row data; keep any columns that the first
|
||||
// statement populated so the schema survives.
|
||||
res.Rows = []map[string]any{{"content": "No record in the database!"}}
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// exesqlQueryOne runs a single statement and returns columns + rows.
|
||||
// Rows are returned as column→value maps with `time.Time` flattened
|
||||
// to "YYYY-MM-DD" (Python pandas `.dt.strftime('%Y-%m-%d')` parity).
|
||||
func exesqlQueryOne(ctx context.Context, db *sql.DB, stmt string, maxRows int) ([]string, []map[string]any, error) {
|
||||
rows, err := db.QueryContext(ctx, stmt)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
out := make([]map[string]any, 0, 16)
|
||||
for rows.Next() {
|
||||
if len(out) >= maxRows {
|
||||
break
|
||||
}
|
||||
raw := make([]any, len(cols))
|
||||
ptrs := make([]any, len(cols))
|
||||
for i := range raw {
|
||||
ptrs[i] = &raw[i]
|
||||
}
|
||||
if err := rows.Scan(ptrs...); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
row := make(map[string]any, len(cols))
|
||||
for i, c := range cols {
|
||||
row[c] = exesqlNormalizeCell(raw[i])
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return cols, out, nil
|
||||
}
|
||||
|
||||
// exesqlNormalizeCell mirrors the Python tool's per-cell conversions:
|
||||
// - time.Time -> "YYYY-MM-DD" (Python: df[col].dt.strftime('%Y-%m-%d'))
|
||||
// - []byte -> string (most drivers return text columns as []byte;
|
||||
// decoding to string gives JSON-friendly output and matches the
|
||||
// Python "df.to_dict(orient='records')" serialization).
|
||||
// - nil / NaN / Inf -> JSON null (Python: df.where(pd.notnull(df), None))
|
||||
// - everything else passes through unchanged.
|
||||
func exesqlNormalizeCell(v any) any {
|
||||
switch x := v.(type) {
|
||||
case nil:
|
||||
return nil
|
||||
case time.Time:
|
||||
return x.Format("2006-01-02")
|
||||
case []byte:
|
||||
return string(x)
|
||||
case float64:
|
||||
// JSON does not represent NaN / Inf; the Python tool drops
|
||||
// them via `convert_decimals` -> None.
|
||||
if isBadFloat(x) {
|
||||
return nil
|
||||
}
|
||||
return x
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func isBadFloat(f float64) bool {
|
||||
// NaN != NaN; Abs(Inf) > max finite. Avoid importing math.
|
||||
if f != f {
|
||||
return true
|
||||
}
|
||||
if f > 1e308 || f < -1e308 {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// splitSQLStatements splits on `;`, ignoring semicolons inside string
|
||||
// literals and line/block comments. This matches what the Python
|
||||
// tool does with `sqls = sql.split(";")` — a naive split, but safe
|
||||
// enough for read-only statements the LLM is expected to produce.
|
||||
func splitSQLStatements(s string) []string {
|
||||
cleaned := stripSQLStrings(stripSQLComments(s))
|
||||
return strings.Split(cleaned, ";")
|
||||
}
|
||||
|
||||
// stripChunkIDMarkers drops the [ID:123] tokens the RAGFlow chunker
|
||||
// sometimes embeds in SQL strings (`re.sub(r"\[ID:[0-9]+\]", "", ...)`
|
||||
// in the Python tool).
|
||||
var exesqlChunkIDRe = regexp.MustCompile(`\[ID:[0-9]+\]`)
|
||||
|
||||
func stripChunkIDMarkers(s string) string { return exesqlChunkIDRe.ReplaceAllString(s, "") }
|
||||
|
||||
// exesqlErrorResult returns the JSON envelope for an error path.
|
||||
func exesqlErrorResult(err error) string {
|
||||
b, _ := json.Marshal(exesqlResult{Error: err.Error()})
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func exesqlMarshalResult(r *exesqlResult) (string, error) {
|
||||
b, err := json.Marshal(r)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("exesql: marshal result: %w", err)
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
// exesqlDriverAndDSN maps a (db_type, conn) tuple to the registered
|
||||
// driver name and DSN. OceanBase reuses the MySQL driver with a
|
||||
// utf8mb4 charset — same trick the Python tool pulls in
|
||||
// `pymysql.connect(..., charset='utf8mb4')`.
|
||||
func exesqlDriverAndDSN(c exesqlConnParams) (driver, dsn string, err error) {
|
||||
switch strings.ToLower(c.DBType) {
|
||||
case "mysql", "mariadb":
|
||||
return "mysql", fmt.Sprintf(
|
||||
"%s:%s@tcp(%s:%d)/%s?parseTime=true&charset=utf8mb4",
|
||||
c.Username, c.Password, c.Host, c.Port, c.Database,
|
||||
), nil
|
||||
case "oceanbase":
|
||||
// OceanBase MySQL-compat mode: same driver, MySQL wire protocol.
|
||||
return "mysql", fmt.Sprintf(
|
||||
"%s:%s@tcp(%s:%d)/%s?parseTime=true&charset=utf8mb4",
|
||||
c.Username, c.Password, c.Host, c.Port, c.Database,
|
||||
), nil
|
||||
case "postgres", "postgresql":
|
||||
return "postgres", fmt.Sprintf(
|
||||
"host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
|
||||
c.Host, c.Port, c.Username, c.Password, c.Database,
|
||||
), nil
|
||||
case "mssql", "sqlserver":
|
||||
return "sqlserver", fmt.Sprintf(
|
||||
"server=%s;port=%d;user id=%s;password=%s;database=%s",
|
||||
c.Host, c.Port, c.Username, c.Password, c.Database,
|
||||
), nil
|
||||
case "trino":
|
||||
return "", "", fmt.Errorf("%w: trino", ErrExeSQLUnsupportedDB)
|
||||
case "ibm db2":
|
||||
return "", "", fmt.Errorf("%w: ibm db2", ErrExeSQLUnsupportedDB)
|
||||
default:
|
||||
return "", "", fmt.Errorf("ExeSQL: unknown db_type %q", c.DBType)
|
||||
}
|
||||
}
|
||||
|
||||
// check returns ErrExeSQLNoCredentials when the required fields are
|
||||
// missing. Mirrors the Python ExeSQLParam.check() but stripped of
|
||||
// the UI-specific "empty value" messages.
|
||||
func (c exesqlConnParams) check() error {
|
||||
if c.DBType == "" || c.Host == "" || c.Username == "" || c.Database == "" {
|
||||
return ErrExeSQLNoCredentials
|
||||
}
|
||||
if c.Port <= 0 {
|
||||
return fmt.Errorf("ExeSQL: invalid port %d", c.Port)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SELECT-only safety validator
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// leadingKeywordRe matches the first non-comment, non-whitespace keyword
|
||||
// in a SQL statement. Comments (-- line, /* block */) and string literals
|
||||
// are stripped before the match.
|
||||
var leadingKeywordRe = regexp.MustCompile(`^[\s,;(]*([A-Za-z]+)`)
|
||||
|
||||
// nonSelectKeywords lists DML/DDL/DCL verbs the parser rejects. These
|
||||
// are the only top-level forms we refuse; everything else (WITH ... SELECT,
|
||||
// SELECT INTO, SHOW, DESCRIBE, EXPLAIN) is allowed because they're
|
||||
// read-only.
|
||||
var nonSelectKeywords = map[string]struct{}{
|
||||
"INSERT": {}, "UPDATE": {}, "DELETE": {}, "REPLACE": {},
|
||||
"TRUNCATE": {},
|
||||
"CREATE": {}, "DROP": {}, "ALTER": {}, "RENAME": {},
|
||||
"GRANT": {}, "REVOKE": {},
|
||||
"LOCK": {}, "UNLOCK": {},
|
||||
"CALL": {}, "EXEC": {}, "EXECUTE": {},
|
||||
"COPY": {},
|
||||
"VACUUM": {}, "ANALYZE": {},
|
||||
"SET": {}, "RESET": {},
|
||||
"USE": {},
|
||||
"KILL": {},
|
||||
"LOAD": {},
|
||||
"CHECKPOINT": {},
|
||||
"BEGIN": {}, "COMMIT": {}, "ROLLBACK": {}, "START": {},
|
||||
"SHUTDOWN": {},
|
||||
}
|
||||
|
||||
// isSelectStatement returns true iff sql is a read-only statement. The
|
||||
// heuristic is intentionally narrow: strip line + block comments and
|
||||
// string literals, scan the first keyword, and reject if it's a
|
||||
// DML/DDL/DCL verb. SQL parsers in Go stdlib don't exist; this matches
|
||||
// the safety bar the Go shell needs in Phase 3 batch 1.
|
||||
func isSelectStatement(sql string) bool {
|
||||
cleaned := stripSQLComments(sql)
|
||||
cleaned = stripSQLStrings(cleaned)
|
||||
m := leadingKeywordRe.FindStringSubmatch(cleaned)
|
||||
if len(m) < 2 {
|
||||
return false
|
||||
}
|
||||
kw := strings.ToUpper(m[1])
|
||||
if _, bad := nonSelectKeywords[kw]; bad {
|
||||
return false
|
||||
}
|
||||
switch kw {
|
||||
case "SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN", "PRAGMA":
|
||||
return true
|
||||
}
|
||||
// Unknown verb → conservative reject. The Python tool would forward
|
||||
// this; the Go shell declines to execute without a recognized form.
|
||||
return false
|
||||
}
|
||||
|
||||
// stripSQLComments removes -- line comments and /* ... */ block comments.
|
||||
// We don't try to handle nested comments (MySQL/PG/SQLite differ) — this
|
||||
// is a best-effort guard for the SELECT validator, not a SQL parser.
|
||||
func stripSQLComments(s string) string {
|
||||
var b strings.Builder
|
||||
b.Grow(len(s))
|
||||
i := 0
|
||||
for i < len(s) {
|
||||
if i+1 < len(s) && s[i] == '-' && s[i+1] == '-' {
|
||||
for i < len(s) && s[i] != '\n' {
|
||||
i++
|
||||
}
|
||||
continue
|
||||
}
|
||||
if i+1 < len(s) && s[i] == '/' && s[i+1] == '*' {
|
||||
i += 2
|
||||
for i+1 < len(s) && !(s[i] == '*' && s[i+1] == '/') {
|
||||
i++
|
||||
}
|
||||
i += 2
|
||||
continue
|
||||
}
|
||||
b.WriteByte(s[i])
|
||||
i++
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// stripSQLStrings removes single- and double-quoted string literals and
|
||||
// replaces them with empty placeholders so that keywords inside string
|
||||
// contents don't confuse the validator.
|
||||
func stripSQLStrings(s string) string {
|
||||
var b strings.Builder
|
||||
b.Grow(len(s))
|
||||
i := 0
|
||||
inStr := byte(0)
|
||||
for i < len(s) {
|
||||
c := s[i]
|
||||
if inStr != 0 {
|
||||
if c == inStr {
|
||||
if i+1 < len(s) && s[i+1] == inStr {
|
||||
b.WriteByte(' ')
|
||||
i += 2
|
||||
continue
|
||||
}
|
||||
inStr = 0
|
||||
}
|
||||
b.WriteByte(' ')
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if c == '\'' || c == '"' || c == '`' {
|
||||
inStr = c
|
||||
b.WriteByte(' ')
|
||||
i++
|
||||
continue
|
||||
}
|
||||
b.WriteByte(c)
|
||||
i++
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
609
internal/agent/tool/exesql_test.go
Normal file
609
internal/agent/tool/exesql_test.go
Normal file
@@ -0,0 +1,609 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/flow/agent/react"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// testConn is a fully-populated connection params struct used by
|
||||
// every test that needs a "valid" tool. Tests that want to exercise
|
||||
// the no-credentials path zero it out.
|
||||
func testConn() exesqlConnParams {
|
||||
return exesqlConnParams{
|
||||
DBType: "mysql",
|
||||
Database: "testdb",
|
||||
Username: "u",
|
||||
Host: "h",
|
||||
Port: 3306,
|
||||
Password: "p",
|
||||
MaxRecords: 100,
|
||||
}
|
||||
}
|
||||
|
||||
// sqlmockDialer returns an exesqlDialer that ignores driver/dsn and
|
||||
// returns a sqlmock-backed *sql.DB. Each call gets a fresh mock so
|
||||
// the test can stage expectations before constructing the tool.
|
||||
func sqlmockDialer(t *testing.T) (exesqlDialer, sqlmock.Sqlmock, func()) {
|
||||
t.Helper()
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
|
||||
if err != nil {
|
||||
t.Fatalf("sqlmock.New: %v", err)
|
||||
}
|
||||
d := func(_, _ string) (*sql.DB, error) { return db, nil }
|
||||
return d, mock, func() { _ = db.Close() }
|
||||
}
|
||||
|
||||
func TestExeSQL_NoCredentials(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewExeSQLTool(exesqlConnParams{}).
|
||||
WithExeSQLDialer(func(_, _ string) (*sql.DB, error) {
|
||||
t.Fatal("dialer should not be called when credentials are missing")
|
||||
return nil, nil
|
||||
})
|
||||
_, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
|
||||
if !errors.Is(err, ErrExeSQLNoCredentials) {
|
||||
t.Fatalf("err = %v, want ErrExeSQLNoCredentials", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_RejectsNonSelect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
sql string
|
||||
}{
|
||||
{"insert", `INSERT INTO foo VALUES (1)`},
|
||||
{"update", `UPDATE foo SET a=1`},
|
||||
{"delete", `DELETE FROM foo`},
|
||||
{"drop", `DROP TABLE foo`},
|
||||
{"create", `CREATE TABLE foo (id INT)`},
|
||||
{"alter", `ALTER TABLE foo ADD COLUMN b INT`},
|
||||
{"truncate", `TRUNCATE foo`},
|
||||
{"grant", `GRANT ALL ON foo TO user`},
|
||||
{"begin", `BEGIN`},
|
||||
{"commit", `COMMIT`},
|
||||
{"set", `SET autocommit=0`},
|
||||
{"kill", `KILL 1234`},
|
||||
{"use", `USE rag_flow`},
|
||||
{"uppercase drop", `DROP DATABASE rag_flow`},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := NewExeSQLTool(testConn())
|
||||
_, err := e.InvokableRun(context.Background(),
|
||||
`{"sql":`+jsonString(c.sql)+`}`)
|
||||
if !errors.Is(err, ErrExeSQLNotSelect) {
|
||||
t.Fatalf("err = %v, want ErrExeSQLNotSelect", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_AllowsSelect(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []string{
|
||||
`SELECT 1`,
|
||||
`select * from t`,
|
||||
` SELECT * FROM t WHERE a = 1`,
|
||||
`WITH cte AS (SELECT 1) SELECT * FROM cte`,
|
||||
`SELECT * FROM t INTO OUTFILE '/tmp/x'`,
|
||||
`SHOW TABLES`,
|
||||
`DESCRIBE t`,
|
||||
`EXPLAIN SELECT * FROM t`,
|
||||
`PRAGMA table_info(t)`,
|
||||
// Keywords inside string literals should be ignored.
|
||||
`SELECT 'DROP TABLE x' AS note FROM dual`,
|
||||
// Line comment with DROP keyword.
|
||||
"-- DROP TABLE foo\nSELECT 1",
|
||||
// Block comment.
|
||||
`/* DROP TABLE foo */ SELECT 1`,
|
||||
}
|
||||
for _, sql := range cases {
|
||||
t.Run(sql, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
// Configure a real-looking query so the validator passes
|
||||
// and the tool reaches the dialer; sqlmock will return no
|
||||
// rows, and we accept either "no record" sentinel or a
|
||||
// sqlmock-driven success.
|
||||
dialer, mock, cleanup := sqlmockDialer(t)
|
||||
defer cleanup()
|
||||
mock.ExpectPing()
|
||||
mock.ExpectQuery("SELECT 1").WillReturnRows(
|
||||
sqlmock.NewRows([]string{"1"}),
|
||||
)
|
||||
e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
|
||||
_, err := e.InvokableRun(context.Background(),
|
||||
`{"sql":`+jsonString(sql)+`}`)
|
||||
// Two acceptable outcomes:
|
||||
// 1. SQL is the literal `SELECT 1` and matches the
|
||||
// mock expectation -> success.
|
||||
// 2. SQL is one of the comment/wrapper variants; the
|
||||
// validator passes but the comment-stripped SQL
|
||||
// differs from the staged expectation -> sqlmock
|
||||
// returns an "unexpected call" error, which is
|
||||
// acceptable because what we're testing here is the
|
||||
// SELECT-only filter, not execution.
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrExeSQLNotSelect) {
|
||||
t.Fatalf("validator rejected a permitted SELECT: %v", err)
|
||||
}
|
||||
// sqlmock mismatch is the expected failure for
|
||||
// comment-stripped variants — fine, the validator
|
||||
// itself passed.
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_RejectsEmptySQL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewExeSQLTool(testConn())
|
||||
_, err := e.InvokableRun(context.Background(), `{"sql":""}`)
|
||||
if err == nil || !strings.Contains(err.Error(), "sql") {
|
||||
t.Fatalf("err = %v, want to mention empty sql", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_RejectsEmptyArgs(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewExeSQLTool(testConn())
|
||||
_, err := e.InvokableRun(context.Background(), "")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty args")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewExeSQLTool(testConn())
|
||||
info, err := e.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "execute_sql" {
|
||||
t.Errorf("Name = %q, want execute_sql", info.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dialer, mock, cleanup := sqlmockDialer(t)
|
||||
defer cleanup()
|
||||
// ExeSQL runs the LLM-supplied SQL verbatim via QueryContext; it
|
||||
// does NOT do database/sql arg binding. Stage the expectation
|
||||
// with the literal value, not "?" + WithArgs.
|
||||
mock.ExpectQuery("SELECT id, name FROM t WHERE id = 7").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "name"}).
|
||||
AddRow(7, "alice").
|
||||
AddRow(8, "bob"))
|
||||
|
||||
e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
|
||||
out, err := e.InvokableRun(context.Background(),
|
||||
`{"sql":"SELECT id, name FROM t WHERE id = 7"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
var got exesqlResult
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v\nout=%s", err, out)
|
||||
}
|
||||
if len(got.Columns) != 2 || got.Columns[0] != "id" || got.Columns[1] != "name" {
|
||||
t.Errorf("Columns = %v, want [id name]", got.Columns)
|
||||
}
|
||||
if len(got.Rows) != 2 {
|
||||
t.Fatalf("Rows = %d, want 2", len(got.Rows))
|
||||
}
|
||||
if got.Rows[0]["name"] != "alice" {
|
||||
t.Errorf("Rows[0][name] = %v, want alice", got.Rows[0]["name"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_ExecuteSelect_NoRowsReturnsSentinel(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dialer, mock, cleanup := sqlmockDialer(t)
|
||||
defer cleanup()
|
||||
mock.ExpectPing()
|
||||
mock.ExpectQuery("SELECT 1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"x"}))
|
||||
|
||||
e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
|
||||
out, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
var got exesqlResult
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v\nout=%s", err, out)
|
||||
}
|
||||
// The Python tool's "No record in the database!" sentinel must
|
||||
// survive the port — downstream nodes (VariableAggregator,
|
||||
// Message) match on it.
|
||||
if len(got.Rows) != 1 || got.Rows[0]["content"] != "No record in the database!" {
|
||||
t.Errorf("Rows = %v, want the no-record sentinel", got.Rows)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_ExecuteSelect_PerStatementErrorIsolated(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dialer, mock, cleanup := sqlmockDialer(t)
|
||||
defer cleanup()
|
||||
mock.ExpectPing()
|
||||
// Two statements, second one errors. Python tool keeps the first
|
||||
// result and records the second as a content entry; the Go port
|
||||
// matches.
|
||||
mock.ExpectQuery("SELECT 1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"x"}).AddRow(1))
|
||||
mock.ExpectQuery("SELECT * FROM bogus").
|
||||
WillReturnError(errors.New("syntax error at or near BOGUS"))
|
||||
|
||||
e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
|
||||
out, err := e.InvokableRun(context.Background(),
|
||||
`{"sql":"SELECT 1; SELECT * FROM bogus"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun should not abort on a per-statement error: %v", err)
|
||||
}
|
||||
var got exesqlResult
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v\nout=%s", err, out)
|
||||
}
|
||||
if len(got.Rows) != 2 {
|
||||
t.Fatalf("Rows = %d, want 2 (one success + one error entry)", len(got.Rows))
|
||||
}
|
||||
// SQL mock returns integers as int64; the JSON round-trip through
|
||||
// map[string]any promotes them to float64. Compare as float64.
|
||||
if x, _ := got.Rows[0]["x"].(float64); x != 1 {
|
||||
t.Errorf("Rows[0][x] = %v (%T), want 1 (the surviving first result)", got.Rows[0]["x"], got.Rows[0]["x"])
|
||||
}
|
||||
if c, _ := got.Rows[1]["content"].(string); !strings.Contains(c, "syntax error") {
|
||||
t.Errorf("Rows[1].content = %q, want to surface the second-statement error", c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_ExecuteSelect_NormalizesTimeAndBytes(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dialer, mock, cleanup := sqlmockDialer(t)
|
||||
defer cleanup()
|
||||
mock.ExpectPing()
|
||||
mock.ExpectQuery("SELECT ts, blob_col FROM t").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"ts", "blob_col"}).
|
||||
AddRow("2024-06-12T03:04:05Z", []byte("hello")))
|
||||
|
||||
e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
|
||||
out, err := e.InvokableRun(context.Background(),
|
||||
`{"sql":"SELECT ts, blob_col FROM t"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
var got exesqlResult
|
||||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v\nout=%s", err, out)
|
||||
}
|
||||
if len(got.Rows) != 1 {
|
||||
t.Fatalf("Rows = %d, want 1", len(got.Rows))
|
||||
}
|
||||
// The mock returns a *string* for "ts" because the column driver
|
||||
// type is text. The normalize step is only triggered for actual
|
||||
// time.Time / []byte values, which is what production drivers
|
||||
// produce. Assert blob_col was decoded to a string instead of
|
||||
// staying as []byte.
|
||||
if _, isBytes := got.Rows[0]["blob_col"].([]byte); isBytes {
|
||||
t.Error("blob_col came back as []byte; exesqlNormalizeCell should convert to string")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_UnsupportedDB(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
e := NewExeSQLTool(exesqlConnParams{
|
||||
DBType: "trino",
|
||||
Host: "h", Port: 8080, Database: "catalog",
|
||||
Username: "u", Password: "p",
|
||||
})
|
||||
_, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
|
||||
if !errors.Is(err, ErrExeSQLUnsupportedDB) {
|
||||
t.Fatalf("err = %v, want ErrExeSQLUnsupportedDB", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExeSQL_DSN_MySQL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
dbType string
|
||||
driver string
|
||||
want string
|
||||
}{
|
||||
{"mysql", "mysql", `u:p@tcp(h:3306)/d?parseTime=true&charset=utf8mb4`},
|
||||
{"mariadb", "mysql", `u:p@tcp(h:3306)/d?parseTime=true&charset=utf8mb4`},
|
||||
{"oceanbase", "mysql", `u:p@tcp(h:3306)/d?parseTime=true&charset=utf8mb4`},
|
||||
{"postgres", "postgres", `host=h port=5432 user=u password=p dbname=d sslmode=disable`},
|
||||
{"mssql", "sqlserver", `server=h;port=1433;user id=u;password=p;database=d`},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.dbType, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
conn := exesqlConnParams{
|
||||
DBType: c.dbType, Host: "h", Port: pickPort(c.dbType),
|
||||
Username: "u", Password: "p", Database: "d",
|
||||
}
|
||||
driver, dsn, err := exesqlDriverAndDSN(conn)
|
||||
if err != nil {
|
||||
t.Fatalf("err: %v", err)
|
||||
}
|
||||
if driver != c.driver {
|
||||
t.Errorf("driver = %q, want %q", driver, c.driver)
|
||||
}
|
||||
if dsn != c.want {
|
||||
t.Errorf("dsn = %q, want %q", dsn, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func pickPort(dbType string) int {
|
||||
switch dbType {
|
||||
case "postgres":
|
||||
return 5432
|
||||
case "mssql":
|
||||
return 1433
|
||||
default:
|
||||
return 3306
|
||||
}
|
||||
}
|
||||
|
||||
// jsonString is a tiny helper to produce a valid JSON string literal
|
||||
// for the SQL values we feed into the tool.
|
||||
func jsonString(s string) string {
|
||||
var b strings.Builder
|
||||
b.WriteByte('"')
|
||||
for _, r := range s {
|
||||
switch r {
|
||||
case '"':
|
||||
b.WriteString(`\"`)
|
||||
case '\\':
|
||||
b.WriteString(`\\`)
|
||||
case '\n':
|
||||
b.WriteString(`\n`)
|
||||
case '\r':
|
||||
b.WriteString(`\r`)
|
||||
case '\t':
|
||||
b.WriteString(`\t`)
|
||||
default:
|
||||
b.WriteRune(r)
|
||||
}
|
||||
}
|
||||
b.WriteByte('"')
|
||||
return string(b.String())
|
||||
}
|
||||
|
||||
// reactScriptedModel is the minimum-viable eino ToolCallingChatModel
|
||||
// that drives the real ReAct loop. It returns a tool_call on the
|
||||
// first Generate and a final content message on the second, recording
|
||||
// every input/output pair so tests can assert on what the framework
|
||||
// actually did (e.g. that a ToolMessage carrying the tool's result
|
||||
// appears in round 2's input).
|
||||
type reactScriptedModel struct {
|
||||
turn int
|
||||
rounds [][]*schema.Message
|
||||
boundTools []*schema.ToolInfo
|
||||
toolName string
|
||||
toolArgs string
|
||||
finalContent string
|
||||
}
|
||||
|
||||
func newReactScriptedModel(toolName, toolArgs, finalContent string) *reactScriptedModel {
|
||||
return &reactScriptedModel{toolName: toolName, toolArgs: toolArgs, finalContent: finalContent}
|
||||
}
|
||||
|
||||
func (m *reactScriptedModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
|
||||
m.boundTools = tools
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (m *reactScriptedModel) Generate(_ context.Context, in []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
||||
cp := make([]*schema.Message, len(in))
|
||||
copy(cp, in)
|
||||
m.rounds = append(m.rounds, cp)
|
||||
m.turn++
|
||||
if m.turn == 1 {
|
||||
return &schema.Message{
|
||||
Role: schema.Assistant,
|
||||
ToolCalls: []schema.ToolCall{{
|
||||
ID: "call_exe_1",
|
||||
Type: "function",
|
||||
Function: schema.FunctionCall{
|
||||
Name: m.toolName,
|
||||
Arguments: m.toolArgs,
|
||||
},
|
||||
}},
|
||||
}, nil
|
||||
}
|
||||
return &schema.Message{Role: schema.Assistant, Content: m.finalContent}, nil
|
||||
}
|
||||
|
||||
func (m *reactScriptedModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
||||
// The ReAct agent drives Generate; Stream is required to satisfy
|
||||
// the interface but never invoked in this test path.
|
||||
sr, sw := schema.Pipe[*schema.Message](1)
|
||||
sw.Close()
|
||||
return sr, nil
|
||||
}
|
||||
|
||||
// TestExeSQL_RealReactAgent_ExecutesTool drives a real eino
|
||||
// react.NewAgent with the real ExeSQLTool (sqlmock-backed DB) and a
|
||||
// scripted chat model. It proves the tool is actually invoked by the
|
||||
// framework, its JSON result is fed back as a ToolMessage on the next
|
||||
// round, and the model emits a final answer grounded in that result.
|
||||
//
|
||||
// This is end-to-end coverage for the "agent -> tool" wiring that the
|
||||
// per-tool unit tests and the registry resolution tests cannot catch:
|
||||
// here the tool descriptor is bound to the model, the model emits a
|
||||
// tool_call, eino's ToolsNode invokes the real ExeSQLTool.InvokableRun,
|
||||
// and the resulting JSON is passed back as a ToolMessage. Replacing
|
||||
// the model with a hand-rolled stub would skip all of that.
|
||||
func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dialer, mock, cleanup := sqlmockDialer(t)
|
||||
defer cleanup()
|
||||
mock.ExpectPing()
|
||||
mock.ExpectQuery("SELECT 42").WillReturnRows(
|
||||
sqlmock.NewRows([]string{"x"}).AddRow(42),
|
||||
)
|
||||
realTool := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
|
||||
|
||||
mdl := newReactScriptedModel(
|
||||
"execute_sql",
|
||||
`{"sql": "SELECT 42"}`,
|
||||
"the answer is 42",
|
||||
)
|
||||
|
||||
agent, err := react.NewAgent(context.Background(), &react.AgentConfig{
|
||||
ToolCallingModel: mdl,
|
||||
ToolsConfig: compose.ToolsNodeConfig{
|
||||
Tools: []einotool.BaseTool{realTool},
|
||||
},
|
||||
MaxStep: 5,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("react.NewAgent: %v", err)
|
||||
}
|
||||
|
||||
out, err := agent.Generate(context.Background(), []*schema.Message{
|
||||
schema.UserMessage("What is 42?"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("agent.Generate: %v", err)
|
||||
}
|
||||
if got, want := out.Content, "the answer is 42"; got != want {
|
||||
t.Errorf("Content = %q, want %q", got, want)
|
||||
}
|
||||
if mdl.turn != 2 {
|
||||
t.Errorf("Generate called %d times, want 2 (tool_call + final)", mdl.turn)
|
||||
}
|
||||
if len(mdl.boundTools) != 1 || mdl.boundTools[0].Name != "execute_sql" {
|
||||
names := make([]string, 0, len(mdl.boundTools))
|
||||
for _, ti := range mdl.boundTools {
|
||||
names = append(names, ti.Name)
|
||||
}
|
||||
t.Errorf("tools bound to model = %v, want [execute_sql]", names)
|
||||
}
|
||||
if len(mdl.rounds) < 2 {
|
||||
t.Fatalf("only %d rounds captured, want >= 2", len(mdl.rounds))
|
||||
}
|
||||
var sawToolResult bool
|
||||
for _, msg := range mdl.rounds[1] {
|
||||
if msg.Role == schema.Tool && strings.Contains(msg.Content, "42") {
|
||||
sawToolResult = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !sawToolResult {
|
||||
t.Errorf("round 2 input did not contain a ToolMessage carrying the tool result; got %d messages", len(mdl.rounds[1]))
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestExeSQL_RealReactAgent_ToolErrorIsolated verifies the
|
||||
// error-as-content path: when the DB returns an error, ExeSQLTool's
|
||||
// InvokableRun surfaces it as a JSON content row (not a Go error).
|
||||
// The eino framework must wrap that as a ToolMessage and pass it to
|
||||
// the model on round 2 without crashing the ReAct loop, so the model
|
||||
// can ground its final answer in the surfaced error.
|
||||
func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dialer, mock, cleanup := sqlmockDialer(t)
|
||||
defer cleanup()
|
||||
mock.ExpectPing()
|
||||
mock.ExpectQuery("SELECT * FROM bogus").
|
||||
WillReturnError(errors.New("syntax error at or near BOGUS"))
|
||||
|
||||
realTool := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
|
||||
|
||||
mdl := newReactScriptedModel(
|
||||
"execute_sql",
|
||||
`{"sql": "SELECT * FROM bogus"}`,
|
||||
"the query had a syntax error",
|
||||
)
|
||||
|
||||
agent, err := react.NewAgent(context.Background(), &react.AgentConfig{
|
||||
ToolCallingModel: mdl,
|
||||
ToolsConfig: compose.ToolsNodeConfig{
|
||||
Tools: []einotool.BaseTool{realTool},
|
||||
},
|
||||
MaxStep: 5,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("react.NewAgent: %v", err)
|
||||
}
|
||||
|
||||
out, err := agent.Generate(context.Background(), []*schema.Message{
|
||||
schema.UserMessage("Find bogus rows"),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("agent.Generate surfaced a Go error; expected the tool to absorb it as content: %v", err)
|
||||
}
|
||||
if got, want := out.Content, "the query had a syntax error"; got != want {
|
||||
t.Errorf("Content = %q, want %q", got, want)
|
||||
}
|
||||
if mdl.turn != 2 {
|
||||
t.Errorf("Generate called %d times, want 2", mdl.turn)
|
||||
}
|
||||
// The round 2 input must include a ToolMessage carrying the
|
||||
// embedded error text — not a Go error and not an empty content.
|
||||
var sawErrorResult bool
|
||||
for _, msg := range mdl.rounds[1] {
|
||||
if msg.Role == schema.Tool && strings.Contains(msg.Content, "syntax error") {
|
||||
sawErrorResult = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !sawErrorResult {
|
||||
t.Errorf("round 2 input did not contain a ToolMessage with the DB error; got %d messages", len(mdl.rounds[1]))
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
174
internal/agent/tool/github.go
Normal file
174
internal/agent/tool/github.go
Normal file
@@ -0,0 +1,174 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const githubToolName = "github"
|
||||
|
||||
const githubToolDescription = "Search GitHub repositories. Returns items[].{full_name, html_url, description, stargazers_count}."
|
||||
|
||||
// githubParams is the JSON shape the model sends into InvokableRun. token
|
||||
// is optional — anonymous requests succeed but are heavily rate-limited
|
||||
// (60/hr vs 5000/hr with a PAT).
|
||||
type githubParams struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
Token string `json:"token"`
|
||||
}
|
||||
|
||||
// githubResult mirrors one element of the upstream `items` array.
|
||||
type githubResult struct {
|
||||
FullName string `json:"full_name"`
|
||||
HTMLURL string `json:"html_url"`
|
||||
Description string `json:"description"`
|
||||
StargazersCount int `json:"stargazers_count"`
|
||||
}
|
||||
|
||||
// githubResponse is the upstream GitHub Search envelope.
|
||||
type githubResponse struct {
|
||||
Items []githubResult `json:"items"`
|
||||
}
|
||||
|
||||
// githubEnvelope is what the model sees.
|
||||
type githubEnvelope struct {
|
||||
Results []githubResult `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// GitHubTool is the Phase 3 batch 3 implementation of the GitHub
|
||||
// repository search tool (plan §2.11.4 row 6, §5 Phase 3 第 3 批). It
|
||||
// GETs the GitHub Search API via the shared HTTPHelper and returns the
|
||||
// top N repository matches.
|
||||
type GitHubTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewGitHubTool returns a GitHubTool using the default HTTPHelper.
|
||||
func NewGitHubTool() *GitHubTool {
|
||||
return NewGitHubToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewGitHubToolWith returns a GitHubTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewGitHubToolWith(h *HTTPHelper) *GitHubTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &GitHubTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (g *GitHubTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: githubToolName,
|
||||
Desc: githubToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query (GitHub search syntax).",
|
||||
Required: true,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of results to return. Defaults to 5 (max 30 per page).",
|
||||
Required: false,
|
||||
},
|
||||
"token": {
|
||||
Type: schema.String,
|
||||
Desc: "Optional GitHub personal access token. Increases rate limit from 60 to 5000 req/hr.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildGitHubURL constructs the GitHub repository search URL. Centralized
|
||||
// so the test suite can verify URL encoding without spinning up a server.
|
||||
func buildGitHubURL(query string, maxResults int) string {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
if maxResults > 30 {
|
||||
maxResults = 30
|
||||
}
|
||||
q := url.Values{}
|
||||
q.Set("q", query)
|
||||
q.Set("per_page", fmt.Sprintf("%d", maxResults))
|
||||
return "https://api.github.com/search/repositories?" + q.Encode()
|
||||
}
|
||||
|
||||
// InvokableRun performs the GitHub repository search.
|
||||
func (g *GitHubTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p githubParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return githubErrJSON(fmt.Errorf("github: parse arguments: %w", err)),
|
||||
fmt.Errorf("github: parse arguments: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(p.Query) == "" {
|
||||
return githubErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("github: query is required")
|
||||
}
|
||||
|
||||
endpoint := buildGitHubURL(p.Query, p.MaxResults)
|
||||
headers := map[string]string{
|
||||
"Accept": "application/vnd.github+json",
|
||||
}
|
||||
if p.Token != "" {
|
||||
headers["Authorization"] = "Bearer " + p.Token
|
||||
}
|
||||
|
||||
resp, err := g.helper.Do(ctx, http.MethodGet, endpoint, "", "", headers)
|
||||
if err != nil {
|
||||
return githubErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return githubErrJSON(fmt.Errorf("github: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("github: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw githubResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return githubErrJSON(fmt.Errorf("github: decode response: %w", err)),
|
||||
fmt.Errorf("github: decode response: %w", err)
|
||||
}
|
||||
return githubJSON(githubEnvelope{Results: raw.Items}), nil
|
||||
}
|
||||
|
||||
func githubJSON(env githubEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"github: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func githubErrJSON(err error) string {
|
||||
return githubJSON(githubEnvelope{Error: err.Error()})
|
||||
}
|
||||
163
internal/agent/tool/github_test.go
Normal file
163
internal/agent/tool/github_test.go
Normal file
@@ -0,0 +1,163 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGitHub_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
query string
|
||||
max int
|
||||
wantPerPage string
|
||||
wantHost string
|
||||
}{
|
||||
{
|
||||
name: "default",
|
||||
query: "ragflow",
|
||||
max: 0,
|
||||
wantPerPage: "5",
|
||||
wantHost: "api.github.com",
|
||||
},
|
||||
{
|
||||
name: "clamped high",
|
||||
query: "x",
|
||||
max: 99,
|
||||
wantPerPage: "30",
|
||||
wantHost: "api.github.com",
|
||||
},
|
||||
{
|
||||
name: "explicit low",
|
||||
query: "language:go stars:>100",
|
||||
max: 3,
|
||||
wantPerPage: "3",
|
||||
wantHost: "api.github.com",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildGitHubURL(tc.query, tc.max)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != tc.wantHost {
|
||||
t.Errorf("host = %q, want %q", u.Host, tc.wantHost)
|
||||
}
|
||||
if u.Path != "/search/repositories" {
|
||||
t.Errorf("path = %q, want /search/repositories", u.Path)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("q") != tc.query {
|
||||
t.Errorf("q = %q, want %q", q.Get("q"), tc.query)
|
||||
}
|
||||
if q.Get("per_page") != tc.wantPerPage {
|
||||
t.Errorf("per_page = %q, want %q", q.Get("per_page"), tc.wantPerPage)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHub_ParseResponse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var gotAuth string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"items": [
|
||||
{"full_name":"infiniflow/ragflow","html_url":"https://github.com/infiniflow/ragflow","description":"RAG engine","stargazers_count":12000},
|
||||
{"full_name":"foo/bar","html_url":"https://github.com/foo/bar","description":"","stargazers_count":5}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewGitHubToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"ragflow","max_results":5,"token":"ghp_xyz"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env githubEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
}
|
||||
if env.Results[0].FullName != "infiniflow/ragflow" {
|
||||
t.Errorf("Results[0].FullName = %q, want infiniflow/ragflow", env.Results[0].FullName)
|
||||
}
|
||||
if env.Results[0].StargazersCount != 12000 {
|
||||
t.Errorf("Results[0].StargazersCount = %d, want 12000", env.Results[0].StargazersCount)
|
||||
}
|
||||
if env.Results[1].Description != "" {
|
||||
t.Errorf("Results[1].Description = %q, want empty", env.Results[1].Description)
|
||||
}
|
||||
if gotAuth != "Bearer ghp_xyz" {
|
||||
t.Errorf("Authorization = %q, want Bearer ghp_xyz", gotAuth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHub_RequiresQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewGitHubTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{"query":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty query")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "query") {
|
||||
t.Errorf("err = %v, want to mention query", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHub_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewGitHubTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "github" {
|
||||
t.Errorf("Name = %q, want github", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "GitHub") {
|
||||
t.Errorf("Desc = %q, want to mention GitHub", info.Desc)
|
||||
}
|
||||
}
|
||||
179
internal/agent/tool/google.go
Normal file
179
internal/agent/tool/google.go
Normal file
@@ -0,0 +1,179 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const googleToolName = "google"
|
||||
|
||||
const googleToolDescription = "Search the web via Google Programmable Search (CSE). Returns items[].{title, link, snippet}."
|
||||
|
||||
// googleParams is the JSON shape the model sends into InvokableRun. The
|
||||
// api_key (CX search-engine id) and cx are both required by the upstream
|
||||
// API; api_key is the Programmable Search API key and cx is the
|
||||
// search-engine ID.
|
||||
type googleParams struct {
|
||||
APIKey string `json:"api_key"`
|
||||
CX string `json:"cx"`
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
// googleResult mirrors one element of the upstream `items` array.
|
||||
type googleResult struct {
|
||||
Title string `json:"title"`
|
||||
Link string `json:"link"`
|
||||
Snippet string `json:"snippet"`
|
||||
}
|
||||
|
||||
// googleResponse is the upstream Programmable Search envelope. We only
|
||||
// model the fields we care about.
|
||||
type googleResponse struct {
|
||||
Items []googleResult `json:"items"`
|
||||
}
|
||||
|
||||
// googleEnvelope is what the model sees.
|
||||
type googleEnvelope struct {
|
||||
Results []googleResult `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// GoogleTool is the Phase 3 batch 2 implementation of the Google
|
||||
// Programmable Search tool (plan §2.11.4 row 10, §5 Phase 3 第 2 批).
|
||||
// It performs a GET against the CSE endpoint using the shared
|
||||
// HTTPHelper.
|
||||
type GoogleTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewGoogleTool returns a GoogleTool using the default HTTPHelper.
|
||||
func NewGoogleTool() *GoogleTool {
|
||||
return NewGoogleToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewGoogleToolWith returns a GoogleTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewGoogleToolWith(h *HTTPHelper) *GoogleTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &GoogleTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (g *GoogleTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: googleToolName,
|
||||
Desc: googleToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query",
|
||||
Required: true,
|
||||
},
|
||||
"api_key": {
|
||||
Type: schema.String,
|
||||
Desc: "Google Programmable Search API key.",
|
||||
Required: true,
|
||||
},
|
||||
"cx": {
|
||||
Type: schema.String,
|
||||
Desc: "Google Programmable Search engine ID (cx).",
|
||||
Required: true,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of results to return. Defaults to 5 (max 10 per request).",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildGoogleURL constructs the CSE query URL. The Programmable Search
|
||||
// API caps `num` at 10; we clamp to that range to avoid upstream errors.
|
||||
func buildGoogleURL(apiKey, cx, query string, maxResults int) string {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
if maxResults > 10 {
|
||||
maxResults = 10
|
||||
}
|
||||
q := url.Values{}
|
||||
q.Set("key", apiKey)
|
||||
q.Set("cx", cx)
|
||||
q.Set("q", query)
|
||||
q.Set("num", fmt.Sprintf("%d", maxResults))
|
||||
return "https://www.googleapis.com/customsearch/v1?" + q.Encode()
|
||||
}
|
||||
|
||||
// InvokableRun performs the Google Programmable Search.
|
||||
func (g *GoogleTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p googleParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return googleErrJSON(fmt.Errorf("google: parse arguments: %w", err)),
|
||||
fmt.Errorf("google: parse arguments: %w", err)
|
||||
}
|
||||
if p.Query == "" {
|
||||
return googleErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("google: query is required")
|
||||
}
|
||||
if p.APIKey == "" || p.CX == "" {
|
||||
return googleErrJSON(fmt.Errorf("google: api_key and cx are required")),
|
||||
fmt.Errorf("google: api_key and cx are required")
|
||||
}
|
||||
|
||||
endpoint := buildGoogleURL(p.APIKey, p.CX, p.Query, p.MaxResults)
|
||||
resp, err := g.helper.Do(ctx, http.MethodGet, endpoint, "", "", nil)
|
||||
if err != nil {
|
||||
return googleErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return googleErrJSON(fmt.Errorf("google: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("google: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw googleResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return googleErrJSON(fmt.Errorf("google: decode response: %w", err)),
|
||||
fmt.Errorf("google: decode response: %w", err)
|
||||
}
|
||||
return googleJSON(googleEnvelope{Results: raw.Items}), nil
|
||||
}
|
||||
|
||||
func googleJSON(env googleEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"google: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func googleErrJSON(err error) string {
|
||||
return googleJSON(googleEnvelope{Error: err.Error()})
|
||||
}
|
||||
355
internal/agent/tool/google_scholar.go
Normal file
355
internal/agent/tool/google_scholar.go
Normal file
@@ -0,0 +1,355 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"golang.org/x/net/html"
|
||||
)
|
||||
|
||||
const googleScholarToolName = "google_scholar"
|
||||
|
||||
const googleScholarToolDescription = "Search Google Scholar for academic articles. Returns {title, link, snippet, authors, year}."
|
||||
|
||||
// googleScholarParams is the JSON shape the model sends into InvokableRun.
|
||||
type googleScholarParams struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
// googleScholarResult is one row in the parsed result list.
|
||||
type googleScholarResult struct {
|
||||
Title string `json:"title"`
|
||||
Link string `json:"link"`
|
||||
Snippet string `json:"snippet"`
|
||||
Authors string `json:"authors"`
|
||||
Year string `json:"year"`
|
||||
}
|
||||
|
||||
// googleScholarEnvelope is what the model sees.
|
||||
type googleScholarEnvelope struct {
|
||||
Results []googleScholarResult `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// googleScholarEndpoint is the Google Scholar search URL. Exposed as
|
||||
// a package var so tests can substitute a httptest.Server URL.
|
||||
var googleScholarEndpoint = "https://scholar.google.com/scholar"
|
||||
|
||||
// GoogleScholarTool is the Phase 3 batch 3 implementation of the
|
||||
// Google Scholar search tool (plan §2.11.4 row 12, §5 Phase 3 第 3 批).
|
||||
// There is no public Scholar API, so we fetch the search-results
|
||||
// HTML and parse it with golang.org/x/net/html.
|
||||
type GoogleScholarTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewGoogleScholarTool returns a GoogleScholarTool using the default
|
||||
// HTTPHelper.
|
||||
func NewGoogleScholarTool() *GoogleScholarTool {
|
||||
return NewGoogleScholarToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewGoogleScholarToolWith returns a GoogleScholarTool that uses the
|
||||
// provided HTTPHelper. Useful for tests.
|
||||
func NewGoogleScholarToolWith(h *HTTPHelper) *GoogleScholarTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &GoogleScholarTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (g *GoogleScholarTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: googleScholarToolName,
|
||||
Desc: googleScholarToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query.",
|
||||
Required: true,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of results to return. Defaults to 5 (max 20 per page).",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildGoogleScholarURL composes the Scholar query URL. Centralized
|
||||
// for testability.
|
||||
func buildGoogleScholarURL(query string, maxResults int) string {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
if maxResults > 20 {
|
||||
maxResults = 20
|
||||
}
|
||||
q := url.Values{}
|
||||
q.Set("q", query)
|
||||
q.Set("hl", "en")
|
||||
q.Set("num", strconv.Itoa(maxResults))
|
||||
return googleScholarEndpoint + "?" + q.Encode()
|
||||
}
|
||||
|
||||
// InvokableRun performs the Google Scholar search.
|
||||
func (g *GoogleScholarTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p googleScholarParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return googleScholarErrJSON(fmt.Errorf("google_scholar: parse arguments: %w", err)),
|
||||
fmt.Errorf("google_scholar: parse arguments: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(p.Query) == "" {
|
||||
return googleScholarErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("google_scholar: query is required")
|
||||
}
|
||||
|
||||
endpoint := buildGoogleScholarURL(p.Query, p.MaxResults)
|
||||
headers := map[string]string{
|
||||
// Scholar blocks obviously non-browser UAs.
|
||||
"User-Agent": "Mozilla/5.0 (compatible; ragflow/1.0)",
|
||||
"Accept": "text/html,application/xhtml+xml",
|
||||
}
|
||||
|
||||
resp, err := g.helper.Do(ctx, http.MethodGet, endpoint, "", "", headers)
|
||||
if err != nil {
|
||||
return googleScholarErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return googleScholarErrJSON(fmt.Errorf("google_scholar: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("google_scholar: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
results, err := parseGoogleScholarHTML(resp.Body, p.MaxResults)
|
||||
if err != nil {
|
||||
return googleScholarErrJSON(fmt.Errorf("google_scholar: parse html: %w", err)),
|
||||
fmt.Errorf("google_scholar: parse html: %w", err)
|
||||
}
|
||||
return googleScholarJSON(googleScholarEnvelope{Results: results}), nil
|
||||
}
|
||||
|
||||
// parseGoogleScholarHTML walks the Scholar search-results HTML and
|
||||
// extracts the conventional .gs_rt / .gs_a / .gs_rs fields. We
|
||||
// deliberately stay defensive: Scholar's markup changes without
|
||||
// notice, so we tolerate missing fields and silently skip articles
|
||||
// that are missing the title.
|
||||
func parseGoogleScholarHTML(body interface {
|
||||
Read(p []byte) (int, error)
|
||||
}, maxResults int) ([]googleScholarResult, error) {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
doc, err := html.Parse(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var results []googleScholarResult
|
||||
var walk func(*html.Node)
|
||||
walk = func(n *html.Node) {
|
||||
if len(results) >= maxResults {
|
||||
return
|
||||
}
|
||||
if n.Type == html.ElementNode {
|
||||
for _, a := range n.Attr {
|
||||
if a.Key == "class" && strings.Contains(a.Val, "gs_ri") {
|
||||
// gs_ri wraps one Scholar result card
|
||||
if r, ok := extractScholarResult(n); ok {
|
||||
results = append(results, r)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for c := n.FirstChild; c != nil; c = c.NextSibling {
|
||||
walk(c)
|
||||
}
|
||||
}
|
||||
walk(doc)
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// extractScholarResult pulls title/link, snippet, and authors/year
|
||||
// from a single .gs_ri node. Returns ok=false when the title anchor
|
||||
// is missing (e.g. PDF / citation links the search layout omits).
|
||||
func extractScholarResult(card *html.Node) (googleScholarResult, bool) {
|
||||
res := googleScholarResult{}
|
||||
|
||||
// Title + link live inside .gs_rt > a
|
||||
title, link := findFirstAnchorInClassedAncestor(card, "gs_rt")
|
||||
if title == "" {
|
||||
return res, false
|
||||
}
|
||||
res.Title = strings.TrimSpace(title)
|
||||
res.Link = link
|
||||
|
||||
// Authors + year live in .gs_a (a single line)
|
||||
if t := findTextWithClass(card, "gs_a"); t != "" {
|
||||
authors, year := splitScholarAuthorsYear(t)
|
||||
res.Authors = authors
|
||||
res.Year = year
|
||||
}
|
||||
|
||||
// Snippet lives in .gs_rs
|
||||
if t := findTextWithClass(card, "gs_rs"); t != "" {
|
||||
res.Snippet = strings.TrimSpace(t)
|
||||
}
|
||||
|
||||
return res, true
|
||||
}
|
||||
|
||||
// findFirstAnchorInClassedAncestor returns the text and href of the
|
||||
// first <a> descendant of n whose ancestor chain contains an element
|
||||
// with `want` in its class list. The `want` argument lets callers
|
||||
// pin the search to a specific Scholar sub-element (e.g. .gs_rt).
|
||||
func findFirstAnchorInClassedAncestor(n *html.Node, want string) (string, string) {
|
||||
var text, href string
|
||||
var found bool
|
||||
var walk func(*html.Node, bool)
|
||||
walk = func(node *html.Node, inTarget bool) {
|
||||
if found {
|
||||
return
|
||||
}
|
||||
here := inTarget
|
||||
if node.Type == html.ElementNode {
|
||||
for _, a := range node.Attr {
|
||||
if a.Key == "class" && strings.Contains(a.Val, want) {
|
||||
here = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if here && node.Data == "a" {
|
||||
for _, a := range node.Attr {
|
||||
if a.Key == "href" {
|
||||
href = a.Val
|
||||
}
|
||||
}
|
||||
text = collectText(node)
|
||||
found = true
|
||||
return
|
||||
}
|
||||
}
|
||||
for c := node.FirstChild; c != nil; c = c.NextSibling {
|
||||
walk(c, here)
|
||||
}
|
||||
}
|
||||
walk(n, false)
|
||||
return text, href
|
||||
}
|
||||
|
||||
// findTextWithClass returns the concatenated text of the first
|
||||
// descendant element that has `want` in its class list. If the
|
||||
// matched element is empty, the search continues into its subtree.
|
||||
func findTextWithClass(n *html.Node, want string) string {
|
||||
var found string
|
||||
var walk func(*html.Node)
|
||||
walk = func(node *html.Node) {
|
||||
if found != "" {
|
||||
return
|
||||
}
|
||||
if node.Type == html.ElementNode {
|
||||
for _, a := range node.Attr {
|
||||
if a.Key == "class" && strings.Contains(a.Val, want) {
|
||||
found = collectText(node)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
for c := node.FirstChild; c != nil; c = c.NextSibling {
|
||||
walk(c)
|
||||
}
|
||||
}
|
||||
walk(n)
|
||||
return found
|
||||
}
|
||||
|
||||
// collectText concatenates all text nodes under n (trimmed of
|
||||
// surrounding whitespace per node).
|
||||
func collectText(n *html.Node) string {
|
||||
var b strings.Builder
|
||||
var walk func(*html.Node)
|
||||
walk = func(node *html.Node) {
|
||||
if node.Type == html.TextNode {
|
||||
b.WriteString(node.Data)
|
||||
}
|
||||
for c := node.FirstChild; c != nil; c = c.NextSibling {
|
||||
walk(c)
|
||||
}
|
||||
}
|
||||
walk(n)
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// splitScholarAuthorsYear parses the .gs_a line, which has the form
|
||||
// "<authors> - <journal>, <year>" or "<authors> - <year>". We pull
|
||||
// the first 4-digit year out and treat everything before " - " as
|
||||
// the author list. Anything we can't parse is returned verbatim so
|
||||
// the model can still see it.
|
||||
func splitScholarAuthorsYear(line string) (authors, year string) {
|
||||
cleaned := strings.TrimSpace(line)
|
||||
// The hyphen between authors and venue is the unicode dash "-".
|
||||
if head, rest, ok := strings.Cut(cleaned, " - "); ok {
|
||||
authors = strings.TrimSpace(head)
|
||||
venue := strings.TrimSpace(rest)
|
||||
year = firstFourDigitYear(venue)
|
||||
return authors, year
|
||||
}
|
||||
year = firstFourDigitYear(cleaned)
|
||||
return cleaned, year
|
||||
}
|
||||
|
||||
// firstFourDigitYear returns the first 4-digit year in s, or "" if
|
||||
// none is found. Years 1900-2099 are recognized.
|
||||
func firstFourDigitYear(s string) string {
|
||||
for i := 0; i+4 <= len(s); i++ {
|
||||
candidate := s[i : i+4]
|
||||
n, err := strconv.Atoi(candidate)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if n >= 1900 && n <= 2099 {
|
||||
return candidate
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func googleScholarJSON(env googleScholarEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"google_scholar: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func googleScholarErrJSON(err error) string {
|
||||
return googleScholarJSON(googleScholarEnvelope{Error: err.Error()})
|
||||
}
|
||||
188
internal/agent/tool/google_scholar_test.go
Normal file
188
internal/agent/tool/google_scholar_test.go
Normal file
@@ -0,0 +1,188 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const cannedScholarHTML = `<html><body>
|
||||
<div class="gs_r gs_or gs_scl" style="display:none"></div>
|
||||
<div class="gs_r">
|
||||
<div class="gs_ri">
|
||||
<h3 class="gs_rt">
|
||||
<a href="https://example.com/paper1">Attention is all you need</a>
|
||||
</h3>
|
||||
<div class="gs_a">Ashish Vaswani, Noam Shazeer, Niki Parmar - NeurIPS 2017</div>
|
||||
<div class="gs_rs">The dominant sequence transduction models are based on complex recurrent or convolutional neural networks.</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="gs_r">
|
||||
<div class="gs_ri">
|
||||
<h3 class="gs_rt">
|
||||
<a href="https://example.com/paper2">BERT: Pre-training of Deep Bidirectional Transformers</a>
|
||||
</h3>
|
||||
<div class="gs_a">Jacob Devlin, Ming-Wei Chang - NAACL 2019</div>
|
||||
<div class="gs_rs">We introduce a new language representation model called BERT.</div>
|
||||
</div>
|
||||
</div>
|
||||
</body></html>`
|
||||
|
||||
func TestGoogleScholar_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
query string
|
||||
max int
|
||||
wantNum string
|
||||
wantHost string
|
||||
wantQuery string
|
||||
}{
|
||||
{
|
||||
name: "default",
|
||||
query: "transformer",
|
||||
max: 0,
|
||||
wantNum: "5",
|
||||
wantHost: "scholar.google.com",
|
||||
wantQuery: "transformer",
|
||||
},
|
||||
{
|
||||
name: "clamped high",
|
||||
query: "x",
|
||||
max: 99,
|
||||
wantNum: "20",
|
||||
wantHost: "scholar.google.com",
|
||||
wantQuery: "x",
|
||||
},
|
||||
{
|
||||
name: "explicit",
|
||||
query: "a b",
|
||||
max: 3,
|
||||
wantNum: "3",
|
||||
wantHost: "scholar.google.com",
|
||||
wantQuery: "a b",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildGoogleScholarURL(tc.query, tc.max)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != tc.wantHost {
|
||||
t.Errorf("host = %q, want %q", u.Host, tc.wantHost)
|
||||
}
|
||||
if u.Path != "/scholar" {
|
||||
t.Errorf("path = %q, want /scholar", u.Path)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("q") != tc.wantQuery {
|
||||
t.Errorf("q = %q, want %q", q.Get("q"), tc.wantQuery)
|
||||
}
|
||||
if q.Get("num") != tc.wantNum {
|
||||
t.Errorf("num = %q, want %q", q.Get("num"), tc.wantNum)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoogleScholar_ParseResults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=UTF-8")
|
||||
_, _ = w.Write([]byte(cannedScholarHTML))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewGoogleScholarToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"transformer","max_results":5}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env googleScholarEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
}
|
||||
if !strings.Contains(env.Results[0].Title, "Attention") {
|
||||
t.Errorf("Results[0].Title = %q, want to contain Attention", env.Results[0].Title)
|
||||
}
|
||||
if env.Results[0].Link != "https://example.com/paper1" {
|
||||
t.Errorf("Results[0].Link = %q, want https://example.com/paper1", env.Results[0].Link)
|
||||
}
|
||||
if !strings.Contains(env.Results[0].Authors, "Vaswani") {
|
||||
t.Errorf("Results[0].Authors = %q, want to contain Vaswani", env.Results[0].Authors)
|
||||
}
|
||||
if env.Results[0].Year != "2017" {
|
||||
t.Errorf("Results[0].Year = %q, want 2017", env.Results[0].Year)
|
||||
}
|
||||
if !strings.Contains(env.Results[0].Snippet, "sequence transduction") {
|
||||
t.Errorf("Results[0].Snippet = %q, want snippet", env.Results[0].Snippet)
|
||||
}
|
||||
if env.Results[1].Year != "2019" {
|
||||
t.Errorf("Results[1].Year = %q, want 2019", env.Results[1].Year)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoogleScholar_RequiresQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewGoogleScholarTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{"query":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty query")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "query") {
|
||||
t.Errorf("err = %v, want to mention query", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoogleScholar_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewGoogleScholarTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "google_scholar" {
|
||||
t.Errorf("Name = %q, want google_scholar", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Scholar") {
|
||||
t.Errorf("Desc = %q, want to mention Scholar", info.Desc)
|
||||
}
|
||||
}
|
||||
170
internal/agent/tool/google_test.go
Normal file
170
internal/agent/tool/google_test.go
Normal file
@@ -0,0 +1,170 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGoogle_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
apiKey string
|
||||
cx string
|
||||
query string
|
||||
max int
|
||||
wantNum string
|
||||
wantHost string
|
||||
}{
|
||||
{
|
||||
name: "default",
|
||||
apiKey: "KEY",
|
||||
cx: "CXID",
|
||||
query: "ragflow",
|
||||
max: 0,
|
||||
wantNum: "5",
|
||||
wantHost: "www.googleapis.com",
|
||||
},
|
||||
{
|
||||
name: "clamped high",
|
||||
apiKey: "K",
|
||||
cx: "C",
|
||||
query: "x",
|
||||
max: 50,
|
||||
wantNum: "10",
|
||||
wantHost: "www.googleapis.com",
|
||||
},
|
||||
{
|
||||
name: "explicit low",
|
||||
apiKey: "K",
|
||||
cx: "C",
|
||||
query: "x y",
|
||||
max: 3,
|
||||
wantNum: "3",
|
||||
wantHost: "www.googleapis.com",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildGoogleURL(tc.apiKey, tc.cx, tc.query, tc.max)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != tc.wantHost {
|
||||
t.Errorf("host = %q, want %q", u.Host, tc.wantHost)
|
||||
}
|
||||
if u.Path != "/customsearch/v1" {
|
||||
t.Errorf("path = %q, want /customsearch/v1", u.Path)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("key") != tc.apiKey {
|
||||
t.Errorf("key = %q, want %q", q.Get("key"), tc.apiKey)
|
||||
}
|
||||
if q.Get("cx") != tc.cx {
|
||||
t.Errorf("cx = %q, want %q", q.Get("cx"), tc.cx)
|
||||
}
|
||||
if q.Get("q") != tc.query {
|
||||
t.Errorf("q = %q, want %q", q.Get("q"), tc.query)
|
||||
}
|
||||
if q.Get("num") != tc.wantNum {
|
||||
t.Errorf("num = %q, want %q", q.Get("num"), tc.wantNum)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoogle_ParseResults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"items": [
|
||||
{"title":"RAGFlow","link":"https://ragflow.io","snippet":"Open source RAG engine"},
|
||||
{"title":"GitHub","link":"https://github.com/infiniflow/ragflow","snippet":"Source code"}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewGoogleToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"ragflow","api_key":"K","cx":"C","max_results":5}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env googleEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
}
|
||||
if env.Results[0].Title != "RAGFlow" {
|
||||
t.Errorf("Results[0].Title = %q, want RAGFlow", env.Results[0].Title)
|
||||
}
|
||||
if env.Results[1].Link != "https://github.com/infiniflow/ragflow" {
|
||||
t.Errorf("Results[1].Link = %q, want https://github.com/infiniflow/ragflow", env.Results[1].Link)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoogle_RequiresAPIKeyAndCX(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewGoogleTool()
|
||||
_, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"x","api_key":"","cx":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for missing api_key and cx")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "api_key") || !strings.Contains(err.Error(), "cx") {
|
||||
t.Errorf("err = %v, want to mention api_key and cx", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoogle_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewGoogleTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "google" {
|
||||
t.Errorf("Name = %q, want google", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Google") {
|
||||
t.Errorf("Desc = %q, want to mention Google", info.Desc)
|
||||
}
|
||||
}
|
||||
396
internal/agent/tool/http_helper.go
Normal file
396
internal/agent/tool/http_helper.go
Normal file
@@ -0,0 +1,396 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// Package tool implements RAGFlow agent canvas tool adapters in Go.
|
||||
//
|
||||
// See plan: .claude/plans/agent-go-port.md §2.11.4 (Tool 移植矩阵, 21 tools)
|
||||
// and §5 Phase 3 (Tool 库, 3-4 周, 按使用频率分 4 批). Phase 3 batch 1 ships
|
||||
// the HTTP底座 (http_helper.go) and 4 必装 tools: Retrieval, CodeExec, ExeSQL,
|
||||
// Crawler. All tools implement eino's tool.InvokableTool interface and are
|
||||
// intended to be registered with the agent canvas via a factory function
|
||||
// (see .claude/plans/agent-go-port.md §2.11.4 "Tool 关键统一模式").
|
||||
package tool
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/rand/v2"
|
||||
"net"
|
||||
"net/http"
|
||||
neturl "net/url"
|
||||
"time"
|
||||
|
||||
"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
|
||||
)
|
||||
|
||||
// RetryConfig controls the HTTPHelper retry policy. The defaults match the
|
||||
// values fixed in plan §5 Phase 3 (3 attempts, 200ms base, 3s max backoff).
|
||||
// Zero values fall back to the defaults so callers can use HTTPHelper{}
|
||||
// or NewHTTPHelper() interchangeably.
|
||||
type RetryConfig struct {
|
||||
// MaxAttempts is the total number of attempts (including the first one).
|
||||
// Values < 1 are clamped to 1 (no retry). Default 3.
|
||||
MaxAttempts int
|
||||
// BaseBackoff is the initial backoff between attempts. Default 200ms.
|
||||
BaseBackoff time.Duration
|
||||
// MaxBackoff caps the exponential backoff. Default 3s.
|
||||
MaxBackoff time.Duration
|
||||
}
|
||||
|
||||
func (c RetryConfig) withDefaults() RetryConfig {
|
||||
if c.MaxAttempts < 1 {
|
||||
c.MaxAttempts = 3
|
||||
}
|
||||
if c.BaseBackoff <= 0 {
|
||||
c.BaseBackoff = 200 * time.Millisecond
|
||||
}
|
||||
if c.MaxBackoff <= 0 {
|
||||
c.MaxBackoff = 3 * time.Second
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// HTTPHelper is a context-aware HTTP client shared by all Phase 3 HTTP tools
|
||||
// (Crawler, plus the Phase 3 batch 2-4 HTTP tools). It wraps http.Client with
|
||||
// otelhttp.NewTransport for OTel span propagation (plan §2.10), enforces a
|
||||
// 30s default timeout, and retries idempotent failures (5xx + network errors)
|
||||
// per the policy in plan §5 Phase 3.
|
||||
//
|
||||
// HTTPHelper is safe for concurrent use by multiple goroutines.
|
||||
type HTTPHelper struct {
|
||||
// baseTransport is the source-of-truth *http.Transport. Do wraps it
|
||||
// with otelhttp and stores the result in h.client. DoPinned clones it
|
||||
// and installs a pinned dialer (see pinnedDialer) so the connect
|
||||
// goes to a known IP regardless of what the request's URL says.
|
||||
// Tests may mutate baseTransport directly to inject a custom
|
||||
// TLSClientConfig (e.g. an in-memory RootCAs pool) — the change
|
||||
// applies to both Do and DoPinned.
|
||||
baseTransport *http.Transport
|
||||
|
||||
// client is the RoundTripper used by Do. It is an OTel-wrapped view
|
||||
// of baseTransport. WithClient replaces this field; DoPinned always
|
||||
// derives its client from baseTransport regardless, so a custom
|
||||
// client installed via WithClient does NOT participate in
|
||||
// DNS-rebinding pinning. Pinning needs the dialer to be replaced at
|
||||
// the transport layer and that is only possible if we own the
|
||||
// transport — hence the baseTransport/client split.
|
||||
client *http.Client
|
||||
|
||||
retry RetryConfig
|
||||
}
|
||||
|
||||
// NewHTTPHelper returns a HTTPHelper with default retry/timeout settings.
|
||||
func NewHTTPHelper() *HTTPHelper {
|
||||
base := &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
DialContext: (&net.Dialer{
|
||||
Timeout: 30 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
}).DialContext,
|
||||
ForceAttemptHTTP2: true,
|
||||
MaxIdleConns: 100,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
TLSHandshakeTimeout: 10 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
}
|
||||
return &HTTPHelper{
|
||||
baseTransport: base,
|
||||
client: &http.Client{
|
||||
Timeout: 30 * time.Second,
|
||||
Transport: otelhttp.NewTransport(base),
|
||||
},
|
||||
retry: RetryConfig{}.withDefaults(),
|
||||
}
|
||||
}
|
||||
|
||||
// NewHTTPHelperWithRetry returns a HTTPHelper with a custom retry policy.
|
||||
// Pass RetryConfig{} to use defaults.
|
||||
func NewHTTPHelperWithRetry(rc RetryConfig) *HTTPHelper {
|
||||
h := NewHTTPHelper()
|
||||
h.retry = rc.withDefaults()
|
||||
return h
|
||||
}
|
||||
|
||||
// WithClient replaces the underlying http.Client used by Do. Useful for
|
||||
// tests that supply a pre-configured transport (e.g. for OTel recording).
|
||||
//
|
||||
// Note: WithClient affects Do calls only. DoPinned always derives its
|
||||
// client from the helper's internal baseTransport, so a custom client
|
||||
// installed here does NOT participate in DNS-rebinding pinning. To
|
||||
// customise pinning behaviour (RootCAs, etc.) mutate baseTransport
|
||||
// directly — typically only tests do this.
|
||||
func (h *HTTPHelper) WithClient(c *http.Client) *HTTPHelper {
|
||||
if c != nil {
|
||||
h.client = c
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
// Do issues an HTTP request and returns the response. The caller is
|
||||
// responsible for closing the returned response body.
|
||||
//
|
||||
// body and contentType may be empty (e.g. for GET). When body is non-empty
|
||||
// and contentType is empty, "application/octet-stream" is assumed so servers
|
||||
// that sniff the body still behave sensibly.
|
||||
//
|
||||
// Retry policy:
|
||||
// - 5xx responses: retried
|
||||
// - network errors (connection refused, DNS, etc.): retried
|
||||
// - 4xx responses: NOT retried (caller error, won't help to retry)
|
||||
// - 2xx / 3xx: returned as-is
|
||||
//
|
||||
// The context is honored on every attempt; cancellation aborts the loop.
|
||||
func (h *HTTPHelper) Do(
|
||||
ctx context.Context,
|
||||
method, url, body, contentType string,
|
||||
headers map[string]string,
|
||||
) (*http.Response, error) {
|
||||
return h.doRawWithClient(ctx, h.client, method, url, body, contentType, headers)
|
||||
}
|
||||
|
||||
// DoPinned is identical to Do but dials pinnedIP for the host in url,
|
||||
// while leaving the request URL host untouched. Use immediately after a
|
||||
// successful ResolveAndValidate (ssrf.go) to defeat DNS-rebinding: an
|
||||
// attacker cannot swap a public A record for a private one between the
|
||||
// validation lookup and the connect, because the connect goes to the IP
|
||||
// we resolved here and the URL host is never re-resolved by the
|
||||
// transport (Go's *http.Transport only re-resolves when the URL's host
|
||||
// is used in the dial, and even then the pinned dialer intercepts).
|
||||
//
|
||||
// The TLS handshake ServerName is derived from the URL host
|
||||
// (Go's http.Transport populates tls.Config.ServerName from
|
||||
// req.URL.Host before calling DialContext), so SNI sends originalHost
|
||||
// and the cert is verified against originalHost — not against the IP.
|
||||
// This is critical: rewriting u.Host to the IP (the previous version of
|
||||
// this fix) sent the IP as SNI and required an IP SAN in the cert,
|
||||
// which is not what real HTTPS sites have. Pinning must happen at the
|
||||
// transport layer, not by mutating the request URL.
|
||||
//
|
||||
// originalHost and pinnedIP are required; an empty value for either
|
||||
// returns an error so misuse is loud rather than silently falling back
|
||||
// to a non-pinned connection. The URL host must also equal originalHost
|
||||
// (defense in depth — a caller that passes a mismatched URL would
|
||||
// produce broken TLS and we refuse rather than silently deliver it).
|
||||
//
|
||||
// Note on proxies: the pinned transport has Proxy disabled, regardless
|
||||
// of baseTransport.Proxy (e.g. HTTP_PROXY/HTTPS_PROXY env). A proxy
|
||||
// would re-resolve the hostname, re-opening the rebinding window the
|
||||
// SSRF guard just closed — and the pinned dialer would otherwise
|
||||
// rewrite the proxy's own dial target to pinnedIP:proxyPort, breaking
|
||||
// the connection. Operators who want proxied outbound should use Do.
|
||||
func (h *HTTPHelper) DoPinned(
|
||||
ctx context.Context,
|
||||
method, url, body, contentType string,
|
||||
headers map[string]string,
|
||||
originalHost string,
|
||||
pinnedIP net.IP,
|
||||
) (*http.Response, error) {
|
||||
if originalHost == "" || pinnedIP == nil {
|
||||
return nil, errors.New("http_helper: DoPinned requires originalHost and pinnedIP")
|
||||
}
|
||||
u, err := neturl.Parse(url)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http_helper: parse url: %w", err)
|
||||
}
|
||||
if u.Hostname() != originalHost {
|
||||
return nil, fmt.Errorf(
|
||||
"http_helper: DoPinned url host %q != originalHost %q "+
|
||||
"(refusing to pin — would break TLS SNI / cert check)",
|
||||
u.Hostname(), originalHost)
|
||||
}
|
||||
return h.doRawWithClient(ctx, h.pinnedClient(pinnedIP), method, url, body, contentType, headers)
|
||||
}
|
||||
|
||||
// pinnedClient builds a one-shot *http.Client whose transport dials
|
||||
// pinnedIP:port instead of originalHost:port. The transport is cloned
|
||||
// from baseTransport so the rest of the connection behaviour (TLS
|
||||
// config, idle pool) is identical to a non-pinned call.
|
||||
//
|
||||
// The proxy setting is explicitly disabled. Two reasons:
|
||||
//
|
||||
// 1. *http.Transport dials the proxy first, and pinnedDialer would
|
||||
// rewrite the proxy's own address to pinnedIP:proxyPort — that
|
||||
// would redirect the proxy dial to a port on the pinned IP and
|
||||
// fail connection establishment in any proxied deployment.
|
||||
//
|
||||
// 2. Even if the dialer were proxy-aware, the proxy itself receives
|
||||
// the request URL (with the original hostname) and re-resolves
|
||||
// it, which re-opens the DNS-rebinding window the SSRF guard
|
||||
// just closed. The pinned dialer must be the only thing that
|
||||
// decides where the TCP connection goes.
|
||||
//
|
||||
// Operators who want proxied outbound for pinned traffic should use Do
|
||||
// (without pinning) and validate upstream of the proxy themselves.
|
||||
//
|
||||
// Design note: "pinned + proxy" is intentionally not supported today.
|
||||
// A proxy that re-resolves the original hostname would defeat the
|
||||
// pinning; only a transparent proxy (one that forwards by IP and does
|
||||
// not re-resolve) could coexist with SSRF-pinned traffic, and the
|
||||
// Go stdlib's http.Transport.Proxy contract does not give us a
|
||||
// reliable way to express that. If a future deployment needs
|
||||
// pinned traffic through a specific egress proxy, add a dedicated
|
||||
// DoPinnedWithProxy (or an HTTPHelper config option that takes an
|
||||
// already-pinned IP) — do not silently re-enable Proxy here without
|
||||
// also re-deriving the SSRF guarantee for that path.
|
||||
func (h *HTTPHelper) pinnedClient(pinnedIP net.IP) *http.Client {
|
||||
base := h.baseTransport.Clone()
|
||||
base.Proxy = nil
|
||||
base.DialContext = (&pinnedDialer{
|
||||
pinnedIP: pinnedIP,
|
||||
base: &net.Dialer{
|
||||
Timeout: 30 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
},
|
||||
}).DialContext
|
||||
return &http.Client{
|
||||
Timeout: h.client.Timeout,
|
||||
Transport: otelhttp.NewTransport(base),
|
||||
}
|
||||
}
|
||||
|
||||
// pinnedDialer is a net.Dialer-compatible DialContext that rewrites the
|
||||
// destination address to pinnedIP:port. The host part of the dial
|
||||
// address is discarded; only the port is preserved. The host of the
|
||||
// HTTP request is not in scope here — that lives on the request, where
|
||||
// it is used by the TLS layer for SNI / cert verification.
|
||||
type pinnedDialer struct {
|
||||
pinnedIP net.IP
|
||||
base *net.Dialer
|
||||
}
|
||||
|
||||
func (d *pinnedDialer) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
_, port, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http_helper: pinned dial: split %q: %w", addr, err)
|
||||
}
|
||||
return d.base.DialContext(ctx, network, net.JoinHostPort(d.pinnedIP.String(), port))
|
||||
}
|
||||
|
||||
// doRawWithClient is the shared retry loop behind Do and DoPinned. The
|
||||
// client decides the transport (and therefore the dialer); Do uses the
|
||||
// helper's default OTel-wrapped client, DoPinned uses a per-call client
|
||||
// with a transport-level pinned dialer.
|
||||
func (h *HTTPHelper) doRawWithClient(
|
||||
ctx context.Context,
|
||||
client *http.Client,
|
||||
method, url, body, contentType string,
|
||||
headers map[string]string,
|
||||
) (*http.Response, error) {
|
||||
if method == "" {
|
||||
method = http.MethodGet
|
||||
}
|
||||
if body != "" && contentType == "" {
|
||||
contentType = "application/octet-stream"
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for attempt := 1; attempt <= h.retry.MaxAttempts; attempt++ {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader([]byte(body)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http_helper: build request: %w", err)
|
||||
}
|
||||
if contentType != "" {
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
}
|
||||
for k, v := range headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
if !isRetryableNetError(err) {
|
||||
return nil, err
|
||||
}
|
||||
if attempt == h.retry.MaxAttempts {
|
||||
return nil, fmt.Errorf("http_helper: %s %s failed after %d attempts: %w", method, SanitizeURL(url), attempt, err)
|
||||
}
|
||||
sleepCtx(ctx, backoff(h.retry.BaseBackoff, h.retry.MaxBackoff, attempt))
|
||||
continue
|
||||
}
|
||||
|
||||
// 5xx is retryable, 4xx is not.
|
||||
if resp.StatusCode >= 500 {
|
||||
lastErr = fmt.Errorf("http_helper: %s %s returned %d", method, SanitizeURL(url), resp.StatusCode)
|
||||
// drain body so the connection can be reused
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if attempt == h.retry.MaxAttempts {
|
||||
return nil, lastErr
|
||||
}
|
||||
sleepCtx(ctx, backoff(h.retry.BaseBackoff, h.retry.MaxBackoff, attempt))
|
||||
continue
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
if lastErr != nil {
|
||||
return nil, lastErr
|
||||
}
|
||||
return nil, errors.New("http_helper: exhausted retries with no recorded error")
|
||||
}
|
||||
|
||||
// backoff returns an exponentially increasing duration with full jitter,
|
||||
// capped at max. attempt is 1-indexed; the first retry uses BaseBackoff.
|
||||
func backoff(base, max time.Duration, attempt int) time.Duration {
|
||||
if attempt < 1 {
|
||||
attempt = 1
|
||||
}
|
||||
d := base
|
||||
for i := 1; i < attempt; i++ {
|
||||
d *= 2
|
||||
if d > max {
|
||||
d = max
|
||||
break
|
||||
}
|
||||
}
|
||||
// full jitter: randomize in [0, d] to avoid thundering herd
|
||||
return time.Duration(rand.Int64N(int64(d) + 1))
|
||||
}
|
||||
|
||||
// sleepCtx waits for d, returning early if ctx is canceled.
|
||||
func sleepCtx(ctx context.Context, d time.Duration) {
|
||||
if d <= 0 {
|
||||
return
|
||||
}
|
||||
t := time.NewTimer(d)
|
||||
defer t.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
|
||||
// isRetryableNetError reports whether a transport-level error should trigger
|
||||
// a retry. Context cancellation / deadline-exceeded are NOT retryable — the
|
||||
// caller explicitly asked for that.
|
||||
func isRetryableNetError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
546
internal/agent/tool/http_helper_test.go
Normal file
546
internal/agent/tool/http_helper_test.go
Normal file
@@ -0,0 +1,546 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"crypto/x509/pkix"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/big"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
neturl "net/url"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func newTestHelper(maxAttempts int, base, max time.Duration) *HTTPHelper {
|
||||
return NewHTTPHelperWithRetry(RetryConfig{
|
||||
MaxAttempts: maxAttempts,
|
||||
BaseBackoff: base,
|
||||
MaxBackoff: max,
|
||||
})
|
||||
}
|
||||
|
||||
// TestHTTPHelper_HappyPath verifies a 2xx response is returned on the first
|
||||
// attempt with no retry, and the body / content-type round-trip cleanly.
|
||||
func TestHTTPHelper_HappyPath(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = io.WriteString(w, `{"ok":true}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond)
|
||||
resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Do returned error: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", resp.StatusCode)
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read body: %v", err)
|
||||
}
|
||||
if string(body) != `{"ok":true}` {
|
||||
t.Fatalf("body = %q, want %q", body, `{"ok":true}`)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPHelper_RetriesOn5xx verifies that the helper retries on 5xx and
|
||||
// returns the first 2xx response. Server returns 503 twice, then 200.
|
||||
func TestHTTPHelper_RetriesOn5xx(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var hits int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
n := atomic.AddInt32(&hits, 1)
|
||||
if n < 3 {
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = io.WriteString(w, "recovered")
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond)
|
||||
resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Do returned error: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if got := atomic.LoadInt32(&hits); got != 3 {
|
||||
t.Fatalf("server hits = %d, want 3 (2 retries + 1 success)", got)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", resp.StatusCode)
|
||||
}
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if string(body) != "recovered" {
|
||||
t.Fatalf("body = %q, want %q", body, "recovered")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPHelper_NoRetryOn4xx verifies that 4xx is returned immediately with
|
||||
// no retry — the caller is responsible for fixing 4xx, retrying won't help.
|
||||
func TestHTTPHelper_NoRetryOn4xx(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var hits int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&hits, 1)
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_, _ = io.WriteString(w, "bad")
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond)
|
||||
resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Do returned error: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if got := atomic.LoadInt32(&hits); got != 1 {
|
||||
t.Fatalf("server hits = %d, want 1 (no retry on 4xx)", got)
|
||||
}
|
||||
if resp.StatusCode != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want 400", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPHelper_Timeout verifies that a context deadline aborts the call
|
||||
// and returns context.DeadlineExceeded promptly (no infinite retry).
|
||||
func TestHTTPHelper_Timeout(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := NewHTTPHelper().WithClient(&http.Client{
|
||||
Timeout: 30 * time.Second,
|
||||
})
|
||||
// Tight deadline — server takes 500ms; we expect to bail in <100ms.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
start := time.Now()
|
||||
_, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil)
|
||||
elapsed := time.Since(start)
|
||||
|
||||
if err == nil {
|
||||
t.Fatal("expected timeout error, got nil")
|
||||
}
|
||||
if elapsed > 250*time.Millisecond {
|
||||
t.Fatalf("Do took %s, want < 250ms (timeout should abort promptly)", elapsed)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPHelper_5xxExhaustion verifies that after MaxAttempts the last
|
||||
// 5xx error is returned.
|
||||
func TestHTTPHelper_5xxExhaustion(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var hits int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&hits, 1)
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond)
|
||||
_, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error after 5xx exhaustion, got nil")
|
||||
}
|
||||
if got := atomic.LoadInt32(&hits); got != 3 {
|
||||
t.Fatalf("server hits = %d, want 3", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPHelper_HeadersAndContentType verifies the helper propagates
|
||||
// custom headers and a non-empty content-type on POST bodies.
|
||||
func TestHTTPHelper_HeadersAndContentType(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if got := r.Header.Get("X-Token"); got != "abc" {
|
||||
t.Errorf("X-Token = %q, want abc", got)
|
||||
}
|
||||
if got := r.Header.Get("Content-Type"); got != "application/json" {
|
||||
t.Errorf("Content-Type = %q, want application/json", got)
|
||||
}
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
w.Header().Set("X-Echo-Body", string(body))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(body)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
h := newTestHelper(1, 1*time.Millisecond, 5*time.Millisecond)
|
||||
resp, err := h.Do(context.Background(), http.MethodPost, srv.URL, `{"k":1}`, "application/json", map[string]string{"X-Token": "abc"})
|
||||
if err != nil {
|
||||
t.Fatalf("Do: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if got := resp.Header.Get("X-Echo-Body"); got != `{"k":1}` {
|
||||
t.Fatalf("echoed body = %q, want %q", got, `{"k":1}`)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBackoffExponential verifies the helper-internal backoff function grows
|
||||
// exponentially and caps at MaxBackoff.
|
||||
func TestBackoffExponential(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
base := 50 * time.Millisecond
|
||||
max := 300 * time.Millisecond
|
||||
|
||||
got1 := backoff(base, max, 1)
|
||||
if got1 < 0 || got1 > base {
|
||||
t.Fatalf("backoff(attempt=1) = %s, want [0, %s]", got1, base)
|
||||
}
|
||||
got3 := backoff(base, max, 3)
|
||||
if got3 < 0 || got3 > max {
|
||||
t.Fatalf("backoff(attempt=3) = %s, want [0, %s] (capped)", got3, max)
|
||||
}
|
||||
// With base=50ms, attempt=10 should be capped at 300ms.
|
||||
got10 := backoff(base, max, 10)
|
||||
if got10 > max {
|
||||
t.Fatalf("backoff(attempt=10) = %s, want <= %s (cap)", got10, max)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRetryConfigDefaults ensures zero-value RetryConfig falls back to
|
||||
// the documented defaults from plan §5 Phase 3.
|
||||
func TestRetryConfigDefaults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := RetryConfig{}.withDefaults()
|
||||
if c.MaxAttempts != 3 {
|
||||
t.Errorf("MaxAttempts = %d, want 3", c.MaxAttempts)
|
||||
}
|
||||
if c.BaseBackoff != 200*time.Millisecond {
|
||||
t.Errorf("BaseBackoff = %s, want 200ms", c.BaseBackoff)
|
||||
}
|
||||
if c.MaxBackoff != 3*time.Second {
|
||||
t.Errorf("MaxBackoff = %s, want 3s", c.MaxBackoff)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPHelper_DoPinnedHTTPS_PreservesSNIAndCert is the regression
|
||||
// test for the M1-rebinding fix as hardened by the post-Phase-7
|
||||
// review: DNS pinning MUST happen at the transport layer, not by
|
||||
// rewriting the request URL. If the URL host were rewritten to the IP,
|
||||
// the TLS ServerName (auto-populated by Go from req.URL.Host) would
|
||||
// become the IP, the SNI would send the IP, and cert verification
|
||||
// would target the IP — which is not what real HTTPS sites have, and
|
||||
// would manifest as x509 errors against any host-cert-only target.
|
||||
//
|
||||
// This test stands up a real TLS server with a cert whose DNS SAN is
|
||||
// "example.test" and whose IP SAN covers the loopback address. The
|
||||
// pinned dialer connects to 127.0.0.1, but the request URL host stays
|
||||
// as "example.test". The server observes the SNI the client sent and
|
||||
// we assert it equals "example.test" (not the IP), and the request
|
||||
// completes successfully (cert verification passes because the URL
|
||||
// host matches the SAN).
|
||||
func TestHTTPHelper_DoPinnedHTTPS_PreservesSNIAndCert(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Cert valid for "example.test" (DNS SAN) and 127.0.0.1, ::1
|
||||
// (IP SANs, just so the test environment itself can resolve).
|
||||
cert := generateTestCert(t, "example.test")
|
||||
|
||||
var observedSNI string
|
||||
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, "ok")
|
||||
}))
|
||||
srv.TLS = &tls.Config{
|
||||
Certificates: []tls.Certificate{cert},
|
||||
GetCertificate: func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
|
||||
observedSNI = hello.ServerName
|
||||
return &cert, nil
|
||||
},
|
||||
}
|
||||
srv.StartTLS()
|
||||
defer srv.Close()
|
||||
|
||||
// srv.URL is https://127.0.0.1:<port>. We extract the port and
|
||||
// re-build the target URL with "example.test" as the host — i.e.
|
||||
// the URL host we send the request with is NOT 127.0.0.1, even
|
||||
// though the connection itself goes to 127.0.0.1.
|
||||
u, perr := neturl.Parse(srv.URL)
|
||||
if perr != nil {
|
||||
t.Fatalf("parse server URL %q: %v", srv.URL, perr)
|
||||
}
|
||||
serverHost, serverPort, sperr := net.SplitHostPort(u.Host)
|
||||
if sperr != nil {
|
||||
t.Fatalf("split server host:port from %q: %v", u.Host, sperr)
|
||||
}
|
||||
targetURL := fmt.Sprintf("https://example.test:%s/", serverPort)
|
||||
pinnedIP := net.ParseIP(serverHost)
|
||||
if pinnedIP == nil {
|
||||
t.Fatalf("server host %q is not an IP literal", serverHost)
|
||||
}
|
||||
|
||||
// Trust the self-signed cert for the duration of the test by
|
||||
// mutating baseTransport directly. This is the supported way to
|
||||
// install a custom trust store (see WithClient's docstring).
|
||||
leaf, lperr := x509.ParseCertificate(cert.Certificate[0])
|
||||
if lperr != nil {
|
||||
t.Fatalf("parse leaf cert: %v", lperr)
|
||||
}
|
||||
pool := x509.NewCertPool()
|
||||
pool.AddCert(leaf)
|
||||
|
||||
h := NewHTTPHelper()
|
||||
h.baseTransport.TLSClientConfig = &tls.Config{RootCAs: pool}
|
||||
|
||||
resp, err := h.DoPinned(context.Background(),
|
||||
http.MethodGet, targetURL, "", "", nil, "example.test", pinnedIP)
|
||||
if err != nil {
|
||||
t.Fatalf("DoPinned: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("status = %d, want 200; body=%s", resp.StatusCode, body)
|
||||
}
|
||||
if observedSNI != "example.test" {
|
||||
t.Errorf("SNI = %q, want %q — DoPinned must keep the URL host for SNI, "+
|
||||
"otherwise HTTPS breaks for any host-cert-only target", observedSNI, "example.test")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPHelper_DoPinnedRefusesMismatchedURLHost locks in the
|
||||
// defense-in-depth check in DoPinned: a caller that passes a URL
|
||||
// whose host does not match originalHost would produce a request
|
||||
// whose TLS ServerName differs from the validated hostname, which is
|
||||
// exactly the SSRF / rebinding bypass we are trying to prevent. We
|
||||
// refuse rather than silently deliver a broken connection.
|
||||
func TestHTTPHelper_DoPinnedRefusesMismatchedURLHost(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
h := NewHTTPHelper()
|
||||
_, err := h.DoPinned(context.Background(),
|
||||
http.MethodGet,
|
||||
"https://attacker.example/foo",
|
||||
"", "", nil,
|
||||
"real.example", // originalHost from resolver
|
||||
net.ParseIP("1.2.3.4"))
|
||||
if err == nil {
|
||||
t.Fatal("DoPinned accepted a mismatched URL host, want error")
|
||||
}
|
||||
if got := err.Error(); !contains(got, "would break TLS SNI") {
|
||||
t.Errorf("error %q does not mention TLS SNI breakage", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPHelper_DoPinnedBypassesProxy locks in the proxy bypass
|
||||
// added after the post-Phase-7 review: even when baseTransport.Proxy
|
||||
// is set (e.g. via HTTP_PROXY / HTTPS_PROXY env), DoPinned must NOT
|
||||
// route through the proxy. Two failure modes were possible without
|
||||
// the bypass:
|
||||
//
|
||||
// 1. *http.Transport dials the proxy first; pinnedDialer would
|
||||
// rewrite the proxy's own address to pinnedIP:proxyPort and the
|
||||
// connection would fail in any proxied deployment.
|
||||
//
|
||||
// 2. Even if the dialer were proxy-aware, the proxy would receive
|
||||
// the original hostname and re-resolve it, re-opening the
|
||||
// rebinding window the SSRF guard just closed.
|
||||
//
|
||||
// The test points baseTransport.Proxy at a port that is guaranteed
|
||||
// to be closed (127.0.0.1:1) — if DoPinned used the proxy, the
|
||||
// request would fail with "connection refused"; if it correctly
|
||||
// bypasses the proxy, the direct dial to 127.0.0.1 succeeds.
|
||||
func TestHTTPHelper_DoPinnedBypassesProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cert := generateTestCert(t, "example.test")
|
||||
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, "ok")
|
||||
}))
|
||||
srv.TLS = &tls.Config{Certificates: []tls.Certificate{cert}}
|
||||
srv.StartTLS()
|
||||
defer srv.Close()
|
||||
|
||||
u, perr := neturl.Parse(srv.URL)
|
||||
if perr != nil {
|
||||
t.Fatalf("parse server URL: %v", perr)
|
||||
}
|
||||
serverHost, serverPort, sperr := net.SplitHostPort(u.Host)
|
||||
if sperr != nil {
|
||||
t.Fatalf("split host:port: %v", sperr)
|
||||
}
|
||||
targetURL := fmt.Sprintf("https://example.test:%s/", serverPort)
|
||||
pinnedIP := net.ParseIP(serverHost)
|
||||
if pinnedIP == nil {
|
||||
t.Fatalf("server host %q is not an IP literal", serverHost)
|
||||
}
|
||||
|
||||
leaf, lperr := x509.ParseCertificate(cert.Certificate[0])
|
||||
if lperr != nil {
|
||||
t.Fatalf("parse leaf cert: %v", lperr)
|
||||
}
|
||||
pool := x509.NewCertPool()
|
||||
pool.AddCert(leaf)
|
||||
|
||||
h := NewHTTPHelper()
|
||||
h.baseTransport.TLSClientConfig = &tls.Config{RootCAs: pool}
|
||||
|
||||
// Point the proxy at a port the OS just gave us and which we then
|
||||
// released. The port is therefore guaranteed to be closed at this
|
||||
// moment (modulo the microsecond race where another process grabs
|
||||
// it between Close() and the test dial — acceptable in a unit
|
||||
// test). This is more self-documenting than reaching for a
|
||||
// "guaranteed-unused" magic port like 127.0.0.1:1.
|
||||
probe, lerr := net.Listen("tcp", "127.0.0.1:0")
|
||||
if lerr != nil {
|
||||
t.Fatalf("listen for a free port: %v", lerr)
|
||||
}
|
||||
closedAddr := probe.Addr().String()
|
||||
if cerr := probe.Close(); cerr != nil {
|
||||
t.Fatalf("close probe listener: %v", cerr)
|
||||
}
|
||||
closedProxy, perr := neturl.Parse("http://" + closedAddr)
|
||||
if perr != nil {
|
||||
t.Fatalf("parse closed proxy URL: %v", perr)
|
||||
}
|
||||
h.baseTransport.Proxy = http.ProxyURL(closedProxy)
|
||||
|
||||
resp, err := h.DoPinned(context.Background(),
|
||||
http.MethodGet, targetURL, "", "", nil, "example.test", pinnedIP)
|
||||
if err != nil {
|
||||
t.Fatalf("DoPinned: %v (proxy may not have been bypassed)", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("status = %d, want 200; body=%s", resp.StatusCode, body)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPinnedDialer_RewritesAddress verifies the unit-level behaviour
|
||||
// of pinnedDialer: it discards the host in the dial address and
|
||||
// connects to pinnedIP:port, preserving the port. This is the
|
||||
// transport-layer primitive that DoPinned uses.
|
||||
func TestPinnedDialer_RewritesAddress(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Stand up a TCP server on 127.0.0.1:<random> and capture the
|
||||
// accepted conn to confirm the pinned dialer actually dialed it.
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("listen: %v", err)
|
||||
}
|
||||
defer ln.Close()
|
||||
|
||||
accepted := make(chan struct{}, 1)
|
||||
go func() {
|
||||
c, aerr := ln.Accept()
|
||||
if aerr == nil {
|
||||
_ = c.Close()
|
||||
}
|
||||
accepted <- struct{}{}
|
||||
}()
|
||||
|
||||
d := &pinnedDialer{
|
||||
pinnedIP: net.ParseIP("127.0.0.1"),
|
||||
base: &net.Dialer{Timeout: 2 * time.Second},
|
||||
}
|
||||
// Pass a deliberately misleading host in the addr — the dialer
|
||||
// must ignore it and dial 127.0.0.1:<ln port> instead.
|
||||
misleading := net.JoinHostPort("203.0.113.99", fmt.Sprint(ln.Addr().(*net.TCPAddr).Port))
|
||||
conn, derr := d.DialContext(context.Background(), "tcp", misleading)
|
||||
if derr != nil {
|
||||
t.Fatalf("DialContext: %v", derr)
|
||||
}
|
||||
_ = conn.Close()
|
||||
|
||||
select {
|
||||
case <-accepted:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("pinned dialer did not connect to 127.0.0.1:port (or listener did not accept)")
|
||||
}
|
||||
}
|
||||
|
||||
// contains is a tiny helper to avoid dragging in strings just for one
|
||||
// assertion. (The rest of the file uses strings.Contains.)
|
||||
func contains(s, sub string) bool {
|
||||
for i := 0; i+len(sub) <= len(s); i++ {
|
||||
if s[i:i+len(sub)] == sub {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// generateTestCert builds a self-signed ECDSA cert valid for dnsName
|
||||
// (DNS SAN) and the loopback addresses (IP SANs). The cert is not
|
||||
// trusted by the system pool — tests must inject it into RootCAs to
|
||||
// use it. Returned together with a t.Cleanup-free form so tests can
|
||||
// store the certificate in their tls.Config.Certificates.
|
||||
func generateTestCert(t *testing.T, dnsName string) tls.Certificate {
|
||||
t.Helper()
|
||||
|
||||
priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
if err != nil {
|
||||
t.Fatalf("ecdsa.GenerateKey: %v", err)
|
||||
}
|
||||
template := x509.Certificate{
|
||||
SerialNumber: big.NewInt(1),
|
||||
Subject: pkix.Name{Organization: []string{"RAGFlow Tool Test"}},
|
||||
NotBefore: time.Now().Add(-time.Hour),
|
||||
NotAfter: time.Now().Add(time.Hour),
|
||||
KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
|
||||
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
|
||||
BasicConstraintsValid: true,
|
||||
DNSNames: []string{dnsName},
|
||||
IPAddresses: []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")},
|
||||
}
|
||||
der, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
|
||||
if err != nil {
|
||||
t.Fatalf("x509.CreateCertificate: %v", err)
|
||||
}
|
||||
return tls.Certificate{
|
||||
Certificate: [][]byte{der},
|
||||
PrivateKey: priv,
|
||||
}
|
||||
}
|
||||
119
internal/agent/tool/jin10.go
Normal file
119
internal/agent/tool/jin10.go
Normal file
@@ -0,0 +1,119 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const jin10ToolName = "jin10"
|
||||
|
||||
const jin10ToolDescription = "Subscribe to Jin10 (jin10.com) real-time financial flash news. " +
|
||||
"STUB: Jin10 has no public REST API; data is delivered over a private " +
|
||||
"WebSocket subscription. Not yet implemented in the Go Canvas. " +
|
||||
"Use the Python Canvas for Jin10 data."
|
||||
|
||||
const jin10UnsupportedMessage = "Jin10 requires WebSocket subscription — not yet implemented in Go Canvas. " +
|
||||
"Use Python Canvas."
|
||||
|
||||
// jin10Params is the JSON shape the model sends into InvokableRun. The
|
||||
// Python implementation accepts a category (e.g. "all", "global", "cny")
|
||||
// and an optional speed filter. The Go stub preserves the shape for
|
||||
// future WebSocket client compatibility but rejects every invocation.
|
||||
type jin10Params struct {
|
||||
Category string `json:"category"`
|
||||
Speed string `json:"speed,omitempty"`
|
||||
}
|
||||
|
||||
// jin10Envelope is the model-facing JSON shape. The stub always
|
||||
// returns a populated Error.
|
||||
type jin10Envelope struct {
|
||||
Items []any `json:"items,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// Jin10Tool is the Phase 3 batch 4 stub implementation of the Jin10
|
||||
// flash-news tool (plan §2.11.4 row 12, §5 Phase 3 第 4 批).
|
||||
//
|
||||
// Jin10 (https://www.jin10.com) has no public API; subscribers receive
|
||||
// data via a private WebSocket connection with token authentication.
|
||||
// A native Go WebSocket client is deferred. For P3-B4 the tool is
|
||||
// registered so DSLs that reference "jin10" keep parsing, but every
|
||||
// invocation fails fast with a clear "use Python Canvas" message.
|
||||
//
|
||||
// Jin10Tool does not own an HTTPHelper — it never makes network calls.
|
||||
type Jin10Tool struct{}
|
||||
|
||||
// NewJin10Tool returns a Jin10Tool. No HTTPHelper is allocated; the
|
||||
// stub never issues network requests.
|
||||
func NewJin10Tool() *Jin10Tool { return &Jin10Tool{} }
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (j *Jin10Tool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: jin10ToolName,
|
||||
Desc: jin10ToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"category": {
|
||||
Type: schema.String,
|
||||
Desc: "News category filter (e.g. \"all\", \"global\", \"cny\"). Defaults to \"all\".",
|
||||
Required: false,
|
||||
},
|
||||
"speed": {
|
||||
Type: schema.String,
|
||||
Desc: "Optional speed/dedup filter. Defaults to empty (no filter).",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// InvokableRun validates the input shape (category is optional with
|
||||
// default "all") and returns a clear "use Python Canvas" error. The
|
||||
// model receives a JSON envelope with the message in the `_ERROR` field.
|
||||
func (j *Jin10Tool) InvokableRun(_ context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p jin10Params
|
||||
if argsJSON != "" {
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return jin10ErrJSON(fmt.Errorf("jin10: parse arguments: %w", err)),
|
||||
errors.New(jin10UnsupportedMessage)
|
||||
}
|
||||
}
|
||||
if p.Category == "" {
|
||||
p.Category = "all"
|
||||
}
|
||||
return jin10ErrJSON(errors.New(jin10UnsupportedMessage)),
|
||||
errors.New(jin10UnsupportedMessage)
|
||||
}
|
||||
|
||||
func jin10JSON(env jin10Envelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"jin10: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func jin10ErrJSON(err error) string {
|
||||
return jin10JSON(jin10Envelope{Error: err.Error()})
|
||||
}
|
||||
101
internal/agent/tool/jin10_test.go
Normal file
101
internal/agent/tool/jin10_test.go
Normal file
@@ -0,0 +1,101 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestJin10_StubsUnsupported(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
args string
|
||||
}{
|
||||
{
|
||||
name: "well-formed args",
|
||||
args: `{"category":"global","speed":"fast"}`,
|
||||
},
|
||||
{
|
||||
name: "default category",
|
||||
args: `{}`,
|
||||
},
|
||||
{
|
||||
name: "empty payload string",
|
||||
args: "",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
tool := NewJin10Tool()
|
||||
out, err := tool.InvokableRun(context.Background(), tc.args)
|
||||
if err == nil {
|
||||
t.Fatalf("expected error, got nil (out=%s)", out)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "WebSocket") {
|
||||
t.Errorf("err = %q, want to mention WebSocket", err.Error())
|
||||
}
|
||||
var env jin10Envelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error == "" {
|
||||
t.Errorf("env.Error = empty, want populated")
|
||||
}
|
||||
if !strings.Contains(env.Error, "WebSocket") {
|
||||
t.Errorf("env.Error = %q, want to mention WebSocket", env.Error)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestJin10_RejectsMalformedJSON(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewJin10Tool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{not json`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for malformed JSON, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "WebSocket") {
|
||||
t.Errorf("err = %q, want to mention WebSocket", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestJin10_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewJin10Tool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "jin10" {
|
||||
t.Errorf("Name = %q, want jin10", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Jin10") {
|
||||
t.Errorf("Desc = %q, want to mention Jin10", info.Desc)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "STUB") && !strings.Contains(info.Desc, "Python") {
|
||||
t.Errorf("Desc = %q, want to flag stub status", info.Desc)
|
||||
}
|
||||
}
|
||||
316
internal/agent/tool/pubmed.go
Normal file
316
internal/agent/tool/pubmed.go
Normal file
@@ -0,0 +1,316 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const pubmedToolName = "pubmed"
|
||||
|
||||
const pubmedToolDescription = "Search PubMed via NCBI E-utilities. Returns {pmid, title, authors, journal, year}."
|
||||
|
||||
// pubmedParams is the JSON shape the model sends into InvokableRun.
|
||||
type pubmedParams struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
// pubmedResult is one row in the returned record list.
|
||||
type pubmedResult struct {
|
||||
PMID string `json:"pmid"`
|
||||
Title string `json:"title"`
|
||||
Authors string `json:"authors"`
|
||||
Journal string `json:"journal"`
|
||||
Year string `json:"year"`
|
||||
}
|
||||
|
||||
// pubmedEnvelope is what the model sees.
|
||||
type pubmedEnvelope struct {
|
||||
Results []pubmedResult `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// pubmedUserAgent is the User-Agent that NCBI requires for all
|
||||
// E-utilities requests. Without it, requests are silently dropped
|
||||
// or rate-limited to a single IP.
|
||||
const pubmedUserAgent = "ragflow/1.0"
|
||||
|
||||
// pubmedESearchEndpoint is the E-utilities esearch URL. Exposed as a
|
||||
// package var so tests can substitute a httptest.Server URL.
|
||||
var pubmedESearchEndpoint = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
|
||||
|
||||
// pubmedESummaryEndpoint is the E-utilities esummary URL. Exposed
|
||||
// as a package var so tests can substitute a httptest.Server URL.
|
||||
var pubmedESummaryEndpoint = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
|
||||
|
||||
// PubMedTool is the Phase 3 batch 3 implementation of the PubMed
|
||||
// search tool (plan §2.11.4 row 13, §5 Phase 3 第 3 批). It uses NCBI
|
||||
// E-utilities: esearch returns a list of PMIDs, then esummary fetches
|
||||
// the full records for those PMIDs.
|
||||
type PubMedTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewPubMedTool returns a PubMedTool using the default HTTPHelper.
|
||||
func NewPubMedTool() *PubMedTool {
|
||||
return NewPubMedToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewPubMedToolWith returns a PubMedTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewPubMedToolWith(h *HTTPHelper) *PubMedTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &PubMedTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (p *PubMedTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: pubmedToolName,
|
||||
Desc: pubmedToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "PubMed search query (full PubMed query syntax supported).",
|
||||
Required: true,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of records to return. Defaults to 5 (max 100 per request).",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildPubMedESearchURL composes the esearch URL. Centralized for
|
||||
// testability.
|
||||
func buildPubMedESearchURL(query string, maxResults int) string {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
if maxResults > 100 {
|
||||
maxResults = 100
|
||||
}
|
||||
q := url.Values{}
|
||||
q.Set("db", "pubmed")
|
||||
q.Set("term", query)
|
||||
q.Set("retmax", strconv.Itoa(maxResults))
|
||||
q.Set("retmode", "json")
|
||||
return pubmedESearchEndpoint + "?" + q.Encode()
|
||||
}
|
||||
|
||||
// buildPubMedESummaryURL composes the esummary URL for a list of
|
||||
// PMIDs. Centralized for testability.
|
||||
func buildPubMedESummaryURL(pmids []string) string {
|
||||
q := url.Values{}
|
||||
q.Set("db", "pubmed")
|
||||
q.Set("id", strings.Join(pmids, ","))
|
||||
q.Set("retmode", "json")
|
||||
return pubmedESummaryEndpoint + "?" + q.Encode()
|
||||
}
|
||||
|
||||
// pubmedESearchResponse is the upstream esearch envelope.
|
||||
type pubmedESearchResponse struct {
|
||||
ESearchResult struct {
|
||||
IDList []string `json:"idlist"`
|
||||
} `json:"esearchresult"`
|
||||
}
|
||||
|
||||
// pubmedESummaryAuthor is one author in the esummary author list.
|
||||
type pubmedESummaryAuthor struct {
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// pubmedESummaryArticle is one article in the esummary result map.
|
||||
type pubmedESummaryArticle struct {
|
||||
Title string `json:"title"`
|
||||
Authors []pubmedESummaryAuthor `json:"authors"`
|
||||
FullJournalName string `json:"fulljournalname"`
|
||||
PubDate string `json:"pubdate"`
|
||||
}
|
||||
|
||||
// pubmedESummaryResponse is the upstream esummary envelope. The
|
||||
// `result` value mixes per-PMID article objects with a string-list
|
||||
// `uids` key, so we decode it as a map of RawMessage and then walk
|
||||
// the entries to extract the article objects (skipping the `uids`
|
||||
// array). See decodePubMedESummary.
|
||||
type pubmedESummaryResponse struct {
|
||||
Result map[string]json.RawMessage `json:"result"`
|
||||
}
|
||||
|
||||
// mustReadAll is a tiny helper to read the entire response body. We
|
||||
// keep it private because pubmed.go owns the only two-step HTTP dance.
|
||||
func mustReadAll(r io.Reader) []byte {
|
||||
b, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// decodePubMedESummary parses the raw esummary response and returns
|
||||
// a PMID → article map. The upstream response is a flat object whose
|
||||
// keys are PMIDs (article values) plus a `uids` key (string list).
|
||||
// A single JSON decode into map[string]Article would fail on `uids`
|
||||
// because Go's encoding/json cannot store arrays in a struct-typed
|
||||
// map; the RawMessage indirection sidesteps that.
|
||||
func decodePubMedESummary(body []byte) (map[string]pubmedESummaryArticle, error) {
|
||||
var raw pubmedESummaryResponse
|
||||
if err := json.Unmarshal(body, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]pubmedESummaryArticle, len(raw.Result))
|
||||
for k, v := range raw.Result {
|
||||
// The `uids` key is a JSON array of strings; skip it.
|
||||
if len(v) > 0 && v[0] == '[' {
|
||||
continue
|
||||
}
|
||||
var article pubmedESummaryArticle
|
||||
if err := json.Unmarshal(v, &article); err != nil {
|
||||
continue
|
||||
}
|
||||
out[k] = article
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// InvokableRun performs the two-step PubMed lookup.
|
||||
func (p *PubMedTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var params pubmedParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), ¶ms); err != nil {
|
||||
return pubmedErrJSON(fmt.Errorf("pubmed: parse arguments: %w", err)),
|
||||
fmt.Errorf("pubmed: parse arguments: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(params.Query) == "" {
|
||||
return pubmedErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("pubmed: query is required")
|
||||
}
|
||||
if params.MaxResults <= 0 {
|
||||
params.MaxResults = 5
|
||||
}
|
||||
|
||||
headers := map[string]string{
|
||||
"User-Agent": pubmedUserAgent,
|
||||
"Accept": "application/json",
|
||||
}
|
||||
|
||||
// Step 1: esearch → list of PMIDs
|
||||
searchURL := buildPubMedESearchURL(params.Query, params.MaxResults)
|
||||
searchResp, err := p.helper.Do(ctx, http.MethodGet, searchURL, "", "", headers)
|
||||
if err != nil {
|
||||
return pubmedErrJSON(err), err
|
||||
}
|
||||
var searchBody pubmedESearchResponse
|
||||
if decErr := json.NewDecoder(searchResp.Body).Decode(&searchBody); decErr != nil {
|
||||
_ = searchResp.Body.Close()
|
||||
return pubmedErrJSON(fmt.Errorf("pubmed: decode esearch: %w", decErr)),
|
||||
fmt.Errorf("pubmed: decode esearch: %w", decErr)
|
||||
}
|
||||
_ = searchResp.Body.Close()
|
||||
if searchResp.StatusCode < 200 || searchResp.StatusCode >= 300 {
|
||||
return pubmedErrJSON(fmt.Errorf("pubmed: esearch returned %d", searchResp.StatusCode)),
|
||||
fmt.Errorf("pubmed: esearch returned %d", searchResp.StatusCode)
|
||||
}
|
||||
|
||||
pmids := searchBody.ESearchResult.IDList
|
||||
if len(pmids) == 0 {
|
||||
return pubmedJSON(pubmedEnvelope{Results: []pubmedResult{}}), nil
|
||||
}
|
||||
|
||||
// Step 2: esummary → full records
|
||||
summaryURL := buildPubMedESummaryURL(pmids)
|
||||
summaryResp, err := p.helper.Do(ctx, http.MethodGet, summaryURL, "", "", headers)
|
||||
if err != nil {
|
||||
return pubmedErrJSON(err), err
|
||||
}
|
||||
defer summaryResp.Body.Close()
|
||||
if summaryResp.StatusCode < 200 || summaryResp.StatusCode >= 300 {
|
||||
return pubmedErrJSON(fmt.Errorf("pubmed: esummary returned %d", summaryResp.StatusCode)),
|
||||
fmt.Errorf("pubmed: esummary returned %d", summaryResp.StatusCode)
|
||||
}
|
||||
|
||||
articles, err := decodePubMedESummary(mustReadAll(summaryResp.Body))
|
||||
if err != nil {
|
||||
return pubmedErrJSON(fmt.Errorf("pubmed: parse esummary: %w", err)),
|
||||
fmt.Errorf("pubmed: parse esummary: %w", err)
|
||||
}
|
||||
|
||||
results := make([]pubmedResult, 0, len(pmids))
|
||||
for _, pmid := range pmids {
|
||||
article, ok := articles[pmid]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
results = append(results, pubmedResult{
|
||||
PMID: pmid,
|
||||
Title: strings.TrimSpace(article.Title),
|
||||
Authors: joinAuthorNames(article.Authors),
|
||||
Journal: article.FullJournalName,
|
||||
Year: firstFourDigitYear(article.PubDate),
|
||||
})
|
||||
}
|
||||
return pubmedJSON(pubmedEnvelope{Results: results}), nil
|
||||
}
|
||||
|
||||
// joinAuthorNames joins the first N authors with ", " and adds
|
||||
// "et al." for any beyond N. We use N=3 to mirror the convention
|
||||
// common in academic citation styles.
|
||||
func joinAuthorNames(authors []pubmedESummaryAuthor) string {
|
||||
const cap = 3
|
||||
if len(authors) == 0 {
|
||||
return ""
|
||||
}
|
||||
if len(authors) <= cap {
|
||||
names := make([]string, len(authors))
|
||||
for i, a := range authors {
|
||||
names[i] = a.Name
|
||||
}
|
||||
return strings.Join(names, ", ")
|
||||
}
|
||||
names := make([]string, 0, cap+1)
|
||||
for i := range cap {
|
||||
names = append(names, authors[i].Name)
|
||||
}
|
||||
names = append(names, "et al.")
|
||||
return strings.Join(names, ", ")
|
||||
}
|
||||
|
||||
func pubmedJSON(env pubmedEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"pubmed: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func pubmedErrJSON(err error) string {
|
||||
return pubmedJSON(pubmedEnvelope{Error: err.Error()})
|
||||
}
|
||||
242
internal/agent/tool/pubmed_test.go
Normal file
242
internal/agent/tool/pubmed_test.go
Normal file
@@ -0,0 +1,242 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPubMed_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("esearch", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildPubMedESearchURL("covid vaccine", 7)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != "eutils.ncbi.nlm.nih.gov" {
|
||||
t.Errorf("host = %q, want eutils.ncbi.nlm.nih.gov", u.Host)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("db") != "pubmed" {
|
||||
t.Errorf("db = %q, want pubmed", q.Get("db"))
|
||||
}
|
||||
if q.Get("term") != "covid vaccine" {
|
||||
t.Errorf("term = %q, want covid vaccine", q.Get("term"))
|
||||
}
|
||||
if q.Get("retmax") != "7" {
|
||||
t.Errorf("retmax = %q, want 7", q.Get("retmax"))
|
||||
}
|
||||
if q.Get("retmode") != "json" {
|
||||
t.Errorf("retmode = %q, want json", q.Get("retmode"))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("esummary", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildPubMedESummaryURL([]string{"12345", "67890"})
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("db") != "pubmed" {
|
||||
t.Errorf("db = %q, want pubmed", q.Get("db"))
|
||||
}
|
||||
if q.Get("id") != "12345,67890" {
|
||||
t.Errorf("id = %q, want 12345,67890", q.Get("id"))
|
||||
}
|
||||
if q.Get("retmode") != "json" {
|
||||
t.Errorf("retmode = %q, want json", q.Get("retmode"))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("esearch clamp high", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildPubMedESearchURL("x", 999)
|
||||
u, _ := url.Parse(got)
|
||||
if u.Query().Get("retmax") != "100" {
|
||||
t.Errorf("retmax = %q, want 100 (clamped)", u.Query().Get("retmax"))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestPubMed_ParseESummary(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var esearchHits, esummaryHits int32
|
||||
var lastUA string
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
lastUA = r.Header.Get("User-Agent")
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if strings.Contains(r.URL.Path, "esearch") {
|
||||
atomic.AddInt32(&esearchHits, 1)
|
||||
_, _ = w.Write([]byte(`{
|
||||
"esearchresult": {
|
||||
"idlist": ["11111", "22222"]
|
||||
}
|
||||
}`))
|
||||
return
|
||||
}
|
||||
if strings.Contains(r.URL.Path, "esummary") {
|
||||
atomic.AddInt32(&esummaryHits, 1)
|
||||
_, _ = w.Write([]byte(`{
|
||||
"result": {
|
||||
"uids": ["11111","22222"],
|
||||
"11111": {
|
||||
"title": "Cochrane review of masks",
|
||||
"authors": [{"name":"Smith J"},{"name":"Doe A"}],
|
||||
"fulljournalname": "Cochrane Database Syst Rev",
|
||||
"pubdate": "2020 Nov 1"
|
||||
},
|
||||
"22222": {
|
||||
"title": "Vaccine efficacy meta-analysis",
|
||||
"authors": [{"name":"Alice"},{"name":"Bob"},{"name":"Carol"},{"name":"Dave"}],
|
||||
"fulljournalname": "Lancet",
|
||||
"pubdate": "2021 Mar-Apr"
|
||||
}
|
||||
}
|
||||
}`))
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewPubMedToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"covid","max_results":5}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
if esearchHits != 1 {
|
||||
t.Errorf("esearch calls = %d, want 1", esearchHits)
|
||||
}
|
||||
if esummaryHits != 1 {
|
||||
t.Errorf("esummary calls = %d, want 1", esummaryHits)
|
||||
}
|
||||
if !strings.Contains(lastUA, "ragflow") {
|
||||
t.Errorf("User-Agent = %q, want to contain ragflow", lastUA)
|
||||
}
|
||||
|
||||
var env pubmedEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
}
|
||||
if env.Results[0].PMID != "11111" {
|
||||
t.Errorf("Results[0].PMID = %q, want 11111", env.Results[0].PMID)
|
||||
}
|
||||
if env.Results[0].Title != "Cochrane review of masks" {
|
||||
t.Errorf("Results[0].Title = %q, want Cochrane review of masks", env.Results[0].Title)
|
||||
}
|
||||
if env.Results[0].Authors != "Smith J, Doe A" {
|
||||
t.Errorf("Results[0].Authors = %q, want Smith J, Doe A", env.Results[0].Authors)
|
||||
}
|
||||
if env.Results[0].Journal != "Cochrane Database Syst Rev" {
|
||||
t.Errorf("Results[0].Journal = %q, want Cochrane Database Syst Rev", env.Results[0].Journal)
|
||||
}
|
||||
if env.Results[0].Year != "2020" {
|
||||
t.Errorf("Results[0].Year = %q, want 2020", env.Results[0].Year)
|
||||
}
|
||||
// 4 authors → first 3 + "et al."
|
||||
if !strings.HasSuffix(env.Results[1].Authors, "et al.") {
|
||||
t.Errorf("Results[1].Authors = %q, want to end with et al.", env.Results[1].Authors)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPubMed_EmptyResults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if strings.Contains(r.URL.Path, "esearch") {
|
||||
_, _ = w.Write([]byte(`{"esearchresult":{"idlist":[]}}`))
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewPubMedToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"noresultsfound-zzz-9999","max_results":5}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
var env pubmedEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 0 {
|
||||
t.Errorf("Results len = %d, want 0", len(env.Results))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPubMed_RequiresQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewPubMedTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{"query":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty query")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "query") {
|
||||
t.Errorf("err = %v, want to mention query", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPubMed_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewPubMedTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "pubmed" {
|
||||
t.Errorf("Name = %q, want pubmed", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "PubMed") {
|
||||
t.Errorf("Desc = %q, want to mention PubMed", info.Desc)
|
||||
}
|
||||
}
|
||||
215
internal/agent/tool/qweather.go
Normal file
215
internal/agent/tool/qweather.go
Normal file
@@ -0,0 +1,215 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const qweatherToolName = "qweather"
|
||||
|
||||
const qweatherToolDescription = "Fetch current weather conditions from 和风天气 (QWeather, devapi.qweather.com). " +
|
||||
"Returns {temp, feelsLike, text, windDir, humidity}."
|
||||
|
||||
// qweatherEndpoint is the QWeather v7 "now" endpoint. Exposed as a
|
||||
// package var so tests can substitute a httptest.Server URL.
|
||||
var qweatherEndpoint = "https://devapi.qweather.com/v7/weather/now"
|
||||
|
||||
// qweatherParams is the JSON shape the model sends into InvokableRun.
|
||||
//
|
||||
// - APIKey (required): the QWeather API key (Console → 项目 → 凭据).
|
||||
// - Location (required): the location code, e.g. "101010100" (Beijing)
|
||||
// or a lat,lon string "39.904,116.405".
|
||||
// - Lang (optional): language code for `text` and `windDir`. Defaults
|
||||
// to "zh" per the inbox spec.
|
||||
type qweatherParams struct {
|
||||
APIKey string `json:"api_key"`
|
||||
Location string `json:"location"`
|
||||
Lang string `json:"lang,omitempty"`
|
||||
}
|
||||
|
||||
// qweatherNow is the upstream `now` object.
|
||||
type qweatherNow struct {
|
||||
Temp string `json:"temp"` // "23"
|
||||
FeelsLike string `json:"feelsLike"` // "22"
|
||||
Text string `json:"text"` // "多云"
|
||||
WindDir string `json:"windDir"` // "东南风"
|
||||
Humidity string `json:"humidity"` // "65"
|
||||
}
|
||||
|
||||
// qweatherResponse is the upstream QWeather envelope. We model only
|
||||
// the fields the inbox spec calls out; additional fields from QWeather
|
||||
// (obsTime, precip, pressure, ...) are ignored.
|
||||
type qweatherResponse struct {
|
||||
Code string `json:"code"` // "200" = OK
|
||||
Now qweatherNow `json:"now"`
|
||||
}
|
||||
|
||||
// qweatherEnvelope is the model-facing JSON shape.
|
||||
type qweatherEnvelope struct {
|
||||
Temp string `json:"temp,omitempty"`
|
||||
FeelsLike string `json:"feels_like,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
WindDir string `json:"wind_dir,omitempty"`
|
||||
Humidity string `json:"humidity,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// QWeatherTool is the Phase 3 batch 4 implementation of the 和风天气
|
||||
// (QWeather) current-conditions tool (plan §2.11.4 row 14, §5 Phase 3
|
||||
// 第 4 批). It performs a GET against devapi.qweather.com/v7/weather/now
|
||||
// and returns the parsed now.{temp, feelsLike, text, windDir, humidity}.
|
||||
//
|
||||
// QWeatherTool uses the shared HTTPHelper for retry/timeout/OTel
|
||||
// propagation.
|
||||
type QWeatherTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewQWeatherTool returns a QWeatherTool using the default HTTPHelper.
|
||||
func NewQWeatherTool() *QWeatherTool {
|
||||
return NewQWeatherToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewQWeatherToolWith returns a QWeatherTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewQWeatherToolWith(h *HTTPHelper) *QWeatherTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &QWeatherTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (q *QWeatherTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: qweatherToolName,
|
||||
Desc: qweatherToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"api_key": {
|
||||
Type: schema.String,
|
||||
Desc: "QWeather API key (Console → 项目 → 凭据).",
|
||||
Required: true,
|
||||
},
|
||||
"location": {
|
||||
Type: schema.String,
|
||||
Desc: "Location code (e.g. 101010100 for Beijing) or \"lat,lon\" (e.g. \"39.904,116.405\").",
|
||||
Required: true,
|
||||
},
|
||||
"lang": {
|
||||
Type: schema.String,
|
||||
Desc: "Language for `text` and `windDir`. Defaults to \"zh\".",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildQWeatherURL composes the devapi.qweather.com URL. Centralized
|
||||
// for testability.
|
||||
func buildQWeatherURL(p qweatherParams) string {
|
||||
q := url.Values{}
|
||||
q.Set("location", p.Location)
|
||||
q.Set("key", p.APIKey)
|
||||
if p.Lang == "" {
|
||||
p.Lang = "zh"
|
||||
}
|
||||
q.Set("lang", p.Lang)
|
||||
return qweatherEndpoint + "?" + q.Encode()
|
||||
}
|
||||
|
||||
// InvokableRun performs the QWeather current-conditions GET.
|
||||
func (q *QWeatherTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p qweatherParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return qweatherErrJSON(fmt.Errorf("qweather: parse arguments: %w", err)),
|
||||
fmt.Errorf("qweather: parse arguments: %w", err)
|
||||
}
|
||||
if p.APIKey == "" {
|
||||
return qweatherErrJSON(fmt.Errorf("qweather: api_key is required")),
|
||||
fmt.Errorf("qweather: api_key is required")
|
||||
}
|
||||
if p.Location == "" {
|
||||
return qweatherErrJSON(fmt.Errorf("qweather: location is required")),
|
||||
fmt.Errorf("qweather: location is required")
|
||||
}
|
||||
|
||||
endpoint := buildQWeatherURL(p)
|
||||
headers := map[string]string{
|
||||
"Accept": "application/json",
|
||||
}
|
||||
|
||||
resp, err := q.helper.Do(ctx, http.MethodGet, endpoint, "", "", headers)
|
||||
if err != nil {
|
||||
return qweatherErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return qweatherErrJSON(fmt.Errorf("qweather: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("qweather: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw qweatherResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return qweatherErrJSON(fmt.Errorf("qweather: decode response: %w", err)),
|
||||
fmt.Errorf("qweather: decode response: %w", err)
|
||||
}
|
||||
// QWeather uses "200" string code for OK; anything else is a
|
||||
// business-level error. Code "404" = 城市不存在, "401" = 认证失败,
|
||||
// "429" = 超过访问次数, "402" = 超过访问速度.
|
||||
if raw.Code != "200" {
|
||||
return qweatherErrJSON(fmt.Errorf("qweather: upstream returned code %q", raw.Code)),
|
||||
fmt.Errorf("qweather: upstream returned code %q", raw.Code)
|
||||
}
|
||||
return qweatherJSON(qweatherEnvelope{
|
||||
Temp: raw.Now.Temp,
|
||||
FeelsLike: raw.Now.FeelsLike,
|
||||
Text: raw.Now.Text,
|
||||
WindDir: raw.Now.WindDir,
|
||||
Humidity: raw.Now.Humidity,
|
||||
}), nil
|
||||
}
|
||||
|
||||
func qweatherJSON(env qweatherEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"qweather: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func qweatherErrJSON(err error) string {
|
||||
return qweatherJSON(qweatherEnvelope{Error: err.Error()})
|
||||
}
|
||||
|
||||
// formatQWeatherError joins a non-2xx status with the upstream code so
|
||||
// the model can see both signals. Exposed for testability.
|
||||
func formatQWeatherError(status int, upstreamCode string) string {
|
||||
if strings.TrimSpace(upstreamCode) == "" {
|
||||
return fmt.Sprintf("qweather: upstream returned %d", status)
|
||||
}
|
||||
return fmt.Sprintf("qweather: upstream returned %d (code=%s)", status, upstreamCode)
|
||||
}
|
||||
230
internal/agent/tool/qweather_test.go
Normal file
230
internal/agent/tool/qweather_test.go
Normal file
@@ -0,0 +1,230 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestQWeather_BuildURL(t *testing.T) {
|
||||
// Not t.Parallel(): this test reads the package-level
|
||||
// `qweatherEndpoint` var, which other tests (running in parallel)
|
||||
// temporarily replace with a httptest.Server URL.
|
||||
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
params qweatherParams
|
||||
wantLoc string
|
||||
wantKey string
|
||||
wantLang string
|
||||
wantHost string
|
||||
wantPath string
|
||||
}{
|
||||
{
|
||||
name: "Beijing, default lang",
|
||||
params: qweatherParams{APIKey: "K-abc", Location: "101010100"},
|
||||
wantLoc: "101010100",
|
||||
wantKey: "K-abc",
|
||||
wantLang: "zh",
|
||||
wantHost: "devapi.qweather.com",
|
||||
wantPath: "/v7/weather/now",
|
||||
},
|
||||
{
|
||||
name: "lat,lon with explicit English",
|
||||
params: qweatherParams{APIKey: "K-xyz", Location: "39.904,116.405", Lang: "en"},
|
||||
wantLoc: "39.904,116.405",
|
||||
wantKey: "K-xyz",
|
||||
wantLang: "en",
|
||||
wantHost: "devapi.qweather.com",
|
||||
wantPath: "/v7/weather/now",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildQWeatherURL(tc.params)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != tc.wantHost {
|
||||
t.Errorf("host = %q, want %q", u.Host, tc.wantHost)
|
||||
}
|
||||
if u.Path != tc.wantPath {
|
||||
t.Errorf("path = %q, want %q", u.Path, tc.wantPath)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("location") != tc.wantLoc {
|
||||
t.Errorf("location = %q, want %q", q.Get("location"), tc.wantLoc)
|
||||
}
|
||||
if q.Get("key") != tc.wantKey {
|
||||
t.Errorf("key = %q, want %q", q.Get("key"), tc.wantKey)
|
||||
}
|
||||
if q.Get("lang") != tc.wantLang {
|
||||
t.Errorf("lang = %q, want %q", q.Get("lang"), tc.wantLang)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQWeather_ParseResponse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var gotQuery url.Values
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotQuery = r.URL.Query()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"code": "200",
|
||||
"updateTime": "2024-05-01T12:00+08:00",
|
||||
"now": {
|
||||
"temp": "23",
|
||||
"feelsLike": "22",
|
||||
"text": "多云",
|
||||
"windDir": "东南风",
|
||||
"humidity": "65"
|
||||
}
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
prev := qweatherEndpoint
|
||||
qweatherEndpoint = srv.URL + "/v7/weather/now"
|
||||
t.Cleanup(func() { qweatherEndpoint = prev })
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewQWeatherToolWith(helper)
|
||||
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"api_key":"K-abc","location":"101010100"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if gotQuery.Get("key") != "K-abc" {
|
||||
t.Errorf("server saw key = %q, want K-abc", gotQuery.Get("key"))
|
||||
}
|
||||
if gotQuery.Get("location") != "101010100" {
|
||||
t.Errorf("server saw location = %q, want 101010100", gotQuery.Get("location"))
|
||||
}
|
||||
if gotQuery.Get("lang") != "zh" {
|
||||
t.Errorf("server saw lang = %q, want zh (default)", gotQuery.Get("lang"))
|
||||
}
|
||||
|
||||
var env qweatherEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if env.Temp != "23" {
|
||||
t.Errorf("Temp = %q, want 23", env.Temp)
|
||||
}
|
||||
if env.FeelsLike != "22" {
|
||||
t.Errorf("FeelsLike = %q, want 22", env.FeelsLike)
|
||||
}
|
||||
if env.Text != "多云" {
|
||||
t.Errorf("Text = %q, want 多云", env.Text)
|
||||
}
|
||||
if env.WindDir != "东南风" {
|
||||
t.Errorf("WindDir = %q, want 东南风", env.WindDir)
|
||||
}
|
||||
if env.Humidity != "65" {
|
||||
t.Errorf("Humidity = %q, want 65", env.Humidity)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQWeather_UpstreamBusinessError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"code":"404","msg":"location not found"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
prev := qweatherEndpoint
|
||||
qweatherEndpoint = srv.URL + "/v7/weather/now"
|
||||
t.Cleanup(func() { qweatherEndpoint = prev })
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewQWeatherToolWith(helper)
|
||||
|
||||
_, err := tool.InvokableRun(context.Background(),
|
||||
`{"api_key":"K-abc","location":"000000000"}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for non-200 upstream code, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "404") {
|
||||
t.Errorf("err = %v, want to surface upstream code", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQWeather_RejectsMissingAPIKey(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewQWeatherTool()
|
||||
_, err := tool.InvokableRun(context.Background(),
|
||||
`{"location":"101010100"}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for missing api_key")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "api_key") {
|
||||
t.Errorf("err = %v, want to mention api_key", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQWeather_RejectsMissingLocation(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewQWeatherTool()
|
||||
_, err := tool.InvokableRun(context.Background(),
|
||||
`{"api_key":"K-abc"}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for missing location")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "location") {
|
||||
t.Errorf("err = %v, want to mention location", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQWeather_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewQWeatherTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "qweather" {
|
||||
t.Errorf("Name = %q, want qweather", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "QWeather") && !strings.Contains(info.Desc, "和风") {
|
||||
t.Errorf("Desc = %q, want to mention QWeather", info.Desc)
|
||||
}
|
||||
}
|
||||
175
internal/agent/tool/registry.go
Normal file
175
internal/agent/tool/registry.go
Normal file
@@ -0,0 +1,175 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
)
|
||||
|
||||
// Factory builds a tool instance by DSL / Agent-visible name and
|
||||
// optional node-level configuration. The config map belongs to the
|
||||
// Agent node / DSL, not to the model-emitted function-call args.
|
||||
type Factory func(params map[string]any) (einotool.BaseTool, error)
|
||||
|
||||
var registry = map[string]Factory{
|
||||
"akshare": noConfig("akshare", func() einotool.BaseTool { return NewAkShareTool() }),
|
||||
"arxiv": noConfig("arxiv", func() einotool.BaseTool { return NewArxivTool() }),
|
||||
"code_exec": noConfig("code_exec", func() einotool.BaseTool { return NewCodeExecTool() }),
|
||||
"crawler": noConfig("crawler", func() einotool.BaseTool { return NewCrawlerTool() }),
|
||||
"deepl": noConfig("deepl", func() einotool.BaseTool { return NewDeepLTool() }),
|
||||
"duckduckgo": noConfig("duckduckgo", func() einotool.BaseTool { return NewDuckDuckGoTool() }),
|
||||
"email": noConfig("email", func() einotool.BaseTool { return NewEmailTool() }),
|
||||
"execute_sql": buildExeSQLTool,
|
||||
"exesql": buildExeSQLTool,
|
||||
"github": noConfig("github", func() einotool.BaseTool { return NewGitHubTool() }),
|
||||
"google": noConfig("google", func() einotool.BaseTool { return NewGoogleTool() }),
|
||||
"google_scholar": noConfig("google_scholar", func() einotool.BaseTool { return NewGoogleScholarTool() }),
|
||||
"jin10": noConfig("jin10", func() einotool.BaseTool { return NewJin10Tool() }),
|
||||
"pubmed": noConfig("pubmed", func() einotool.BaseTool { return NewPubMedTool() }),
|
||||
"qweather": noConfig("qweather", func() einotool.BaseTool { return NewQWeatherTool() }),
|
||||
"retrieval": noConfig("retrieval", func() einotool.BaseTool { return NewRetrievalTool() }),
|
||||
"search_my_dateset": noConfig("search_my_dateset", func() einotool.BaseTool { return NewRetrievalTool() }),
|
||||
"searxng": noConfig("searxng", func() einotool.BaseTool { return NewSearXNGTool() }),
|
||||
"tavily": noConfig("tavily", func() einotool.BaseTool { return NewTavilyTool() }),
|
||||
"tushare": noConfig("tushare", func() einotool.BaseTool { return NewTushareTool() }),
|
||||
"wencai": noConfig("wencai", func() einotool.BaseTool { return NewWencaiTool() }),
|
||||
"wikipedia": noConfig("wikipedia", func() einotool.BaseTool { return NewWikipediaTool() }),
|
||||
"yahoo_finance": noConfig("yahoo_finance", func() einotool.BaseTool { return NewYahooFinanceTool() }),
|
||||
}
|
||||
|
||||
func noConfig(name string, fn func() einotool.BaseTool) Factory {
|
||||
return func(params map[string]any) (einotool.BaseTool, error) {
|
||||
if len(params) != 0 {
|
||||
return nil, fmt.Errorf("agent tool: tool %q does not accept node-level params", name)
|
||||
}
|
||||
return fn(), nil
|
||||
}
|
||||
}
|
||||
|
||||
// BuildByName resolves a tool name into an Eino BaseTool.
|
||||
func BuildByName(name string, params map[string]any) (einotool.BaseTool, error) {
|
||||
key := strings.ToLower(strings.TrimSpace(name))
|
||||
if key == "" {
|
||||
return nil, fmt.Errorf("agent tool: empty tool name")
|
||||
}
|
||||
factory, ok := registry[key]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("agent tool: unsupported tool %q", name)
|
||||
}
|
||||
if factory == nil {
|
||||
return nil, fmt.Errorf("agent tool: nil factory for %q", name)
|
||||
}
|
||||
return factory(params)
|
||||
}
|
||||
|
||||
// BuildAll resolves a list of tool names into Eino BaseTool instances.
|
||||
// perToolParams is keyed by the Agent-visible tool name.
|
||||
func BuildAll(names []string, perToolParams map[string]map[string]any) ([]einotool.BaseTool, error) {
|
||||
if len(names) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
tools := make([]einotool.BaseTool, 0, len(names))
|
||||
for _, name := range names {
|
||||
var params map[string]any
|
||||
if perToolParams != nil {
|
||||
params = perToolParams[strings.ToLower(strings.TrimSpace(name))]
|
||||
if params == nil {
|
||||
params = perToolParams[name]
|
||||
}
|
||||
}
|
||||
t, err := BuildByName(name, params)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tools = append(tools, t)
|
||||
}
|
||||
return tools, nil
|
||||
}
|
||||
|
||||
func buildExeSQLTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
conn, err := decodeExeSQLConnParams(params)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return NewExeSQLTool(conn), nil
|
||||
}
|
||||
|
||||
func decodeExeSQLConnParams(params map[string]any) (exesqlConnParams, error) {
|
||||
if len(params) == 0 {
|
||||
return exesqlConnParams{}, fmt.Errorf(
|
||||
"agent tool: execute_sql requires node-level params " +
|
||||
"(db_type/host/port/database/username/password)",
|
||||
)
|
||||
}
|
||||
conn := exesqlConnParams{}
|
||||
if v, ok := stringParam(params, "db_type"); ok {
|
||||
conn.DBType = v
|
||||
}
|
||||
if v, ok := stringParam(params, "database"); ok {
|
||||
conn.Database = v
|
||||
}
|
||||
if v, ok := stringParam(params, "username"); ok {
|
||||
conn.Username = v
|
||||
}
|
||||
if v, ok := stringParam(params, "host"); ok {
|
||||
conn.Host = v
|
||||
}
|
||||
if v, ok := intParam(params, "port"); ok {
|
||||
conn.Port = v
|
||||
}
|
||||
if v, ok := stringParam(params, "password"); ok {
|
||||
conn.Password = v
|
||||
}
|
||||
if v, ok := intParam(params, "max_records"); ok {
|
||||
conn.MaxRecords = v
|
||||
}
|
||||
if err := conn.check(); err != nil {
|
||||
return exesqlConnParams{}, fmt.Errorf("agent tool: execute_sql config: %w", err)
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func stringParam(params map[string]any, key string) (string, bool) {
|
||||
v, ok := params[key]
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
s, ok := v.(string)
|
||||
return s, ok
|
||||
}
|
||||
|
||||
func intParam(params map[string]any, key string) (int, bool) {
|
||||
v, ok := params[key]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x, true
|
||||
case int32:
|
||||
return int(x), true
|
||||
case int64:
|
||||
return int(x), true
|
||||
case float64:
|
||||
return int(x), true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
186
internal/agent/tool/registry_test.go
Normal file
186
internal/agent/tool/registry_test.go
Normal file
@@ -0,0 +1,186 @@
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBuildAll_KnownTools(t *testing.T) {
|
||||
tools, err := BuildAll([]string{"retrieval", "wikipedia"}, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildAll: %v", err)
|
||||
}
|
||||
if len(tools) != 2 {
|
||||
t.Fatalf("len(tools) = %d, want 2", len(tools))
|
||||
}
|
||||
info0, err := tools[0].Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("tools[0].Info: %v", err)
|
||||
}
|
||||
if info0.Name != "search_my_dateset" {
|
||||
t.Errorf("tools[0].Info().Name = %q, want search_my_dateset", info0.Name)
|
||||
}
|
||||
info1, err := tools[1].Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("tools[1].Info: %v", err)
|
||||
}
|
||||
if info1.Name != "wikipedia" {
|
||||
t.Errorf("tools[1].Info().Name = %q, want wikipedia", info1.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAll_UnknownTool(t *testing.T) {
|
||||
_, err := BuildAll([]string{"does_not_exist"}, nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for unknown tool")
|
||||
}
|
||||
if !strings.Contains(err.Error(), `unsupported tool "does_not_exist"`) {
|
||||
t.Fatalf("err = %q, want unsupported tool message", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAll_AllRegisteredTools(t *testing.T) {
|
||||
names := []string{
|
||||
"akshare", "arxiv", "code_exec", "crawler", "deepl", "duckduckgo",
|
||||
"email", "github", "google", "google_scholar", "jin10", "pubmed",
|
||||
"qweather", "retrieval", "searxng", "tavily", "tushare", "wencai",
|
||||
"wikipedia", "yahoo_finance", "execute_sql",
|
||||
}
|
||||
params := map[string]map[string]any{
|
||||
"execute_sql": {
|
||||
"db_type": "mysql",
|
||||
"host": "127.0.0.1",
|
||||
"port": 3306,
|
||||
"database": "demo",
|
||||
"username": "u",
|
||||
"password": "p",
|
||||
"max_records": 10,
|
||||
},
|
||||
}
|
||||
tools, err := BuildAll(names, params)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildAll(all registered): %v", err)
|
||||
}
|
||||
if len(tools) != len(names) {
|
||||
t.Fatalf("len(tools) = %d, want %d", len(tools), len(names))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAll_ExeSQLRequiresNodeParams(t *testing.T) {
|
||||
_, err := BuildAll([]string{"execute_sql"}, nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected execute_sql config error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "execute_sql requires node-level params") {
|
||||
t.Fatalf("err = %q, want execute_sql config error", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolRegistry_SchemasAreComplete sweeps every name the public
|
||||
// registry advertises (including the execute_sql/exesql and
|
||||
// retrieval/search_my_dateset alias pairs), builds the tool, and
|
||||
// asserts that its Info() returns a complete schema — non-empty
|
||||
// Name and Desc, non-nil ParamsOneOf, and a consistent canonical
|
||||
// name across alias entries. Catches drift like "tool renamed but
|
||||
// registry not updated", "param added but schema not updated",
|
||||
// "tool registered with empty description", and "alias points to
|
||||
// the wrong canonical name".
|
||||
func TestToolRegistry_SchemasAreComplete(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Every entry the registry advertises. 23 names, 22 unique
|
||||
// canonical tools (execute_sql == exesql, retrieval ==
|
||||
// search_my_dateset).
|
||||
names := []string{
|
||||
"akshare", "arxiv", "code_exec", "crawler", "deepl", "duckduckgo",
|
||||
"email", "execute_sql", "exesql", "github", "google",
|
||||
"google_scholar", "jin10", "pubmed", "qweather", "retrieval",
|
||||
"search_my_dateset", "searxng", "tavily", "tushare", "wencai",
|
||||
"wikipedia", "yahoo_finance",
|
||||
}
|
||||
params := map[string]map[string]any{
|
||||
"execute_sql": {
|
||||
"db_type": "mysql",
|
||||
"host": "127.0.0.1",
|
||||
"port": 3306,
|
||||
"database": "demo",
|
||||
"username": "u",
|
||||
"password": "p",
|
||||
"max_records": 10,
|
||||
},
|
||||
"exesql": {
|
||||
"db_type": "mysql",
|
||||
"host": "127.0.0.1",
|
||||
"port": 3306,
|
||||
"database": "demo",
|
||||
"username": "u",
|
||||
"password": "p",
|
||||
"max_records": 10,
|
||||
},
|
||||
}
|
||||
tools, err := BuildAll(names, params)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildAll(%d names): %v", len(names), err)
|
||||
}
|
||||
if len(tools) != len(names) {
|
||||
t.Fatalf("BuildAll returned %d tools for %d names", len(tools), len(names))
|
||||
}
|
||||
|
||||
// Schema-level checks per entry.
|
||||
for i, name := range names {
|
||||
info, err := tools[i].Info(context.Background())
|
||||
if err != nil {
|
||||
t.Errorf("tools[%d] (registry name %q).Info: %v", i, name, err)
|
||||
continue
|
||||
}
|
||||
if info.Name == "" {
|
||||
t.Errorf("tools[%d] (registry name %q).Info().Name is empty", i, name)
|
||||
}
|
||||
if info.Desc == "" {
|
||||
t.Errorf("tools[%d] (registry name %q).Info().Desc is empty", i, name)
|
||||
}
|
||||
if info.ParamsOneOf == nil {
|
||||
t.Errorf("tools[%d] (registry name %q).Info().ParamsOneOf is nil", i, name)
|
||||
}
|
||||
}
|
||||
|
||||
// Alias consistency: execute_sql and exesql must surface the
|
||||
// same canonical Info().Name; same for retrieval and
|
||||
// search_my_dateset. A bug here would mean an alias was
|
||||
// accidentally pointed at a different tool.
|
||||
canonicalByAlias := map[string]string{
|
||||
"execute_sql": "execute_sql",
|
||||
"exesql": "execute_sql",
|
||||
"retrieval": "search_my_dateset",
|
||||
"search_my_dateset": "search_my_dateset",
|
||||
}
|
||||
for _, name := range names {
|
||||
canonical, ok := canonicalByAlias[name]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
idx := indexOf(names, name)
|
||||
info, err := tools[idx].Info(context.Background())
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if info.Name != canonical {
|
||||
t.Errorf("registry name %q: Info().Name = %q, want %q (alias must surface canonical name)",
|
||||
name, info.Name, canonical)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// indexOf returns the index of s in xs, or -1 if not present.
|
||||
// Tiny helper to keep the alias loop above free of a slice lookup
|
||||
// closure; the test's names slice is <30 items so linear scan is
|
||||
// fine.
|
||||
func indexOf(xs []string, s string) int {
|
||||
for i, x := range xs {
|
||||
if x == s {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
163
internal/agent/tool/retrieval.go
Normal file
163
internal/agent/tool/retrieval.go
Normal file
@@ -0,0 +1,163 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// ErrGraphRAGNotSupported is returned by the Retrieval tool when callers
|
||||
// pass use_kg=true. GraphRAG is explicitly out of scope for the Go Canvas
|
||||
// (plan §5 Phase 3 + §9 Q3); users must either disable use_kg or fall back
|
||||
// to the Python Canvas.
|
||||
var ErrGraphRAGNotSupported = errors.New("GraphRAG 检索暂不支持,请使用 Python Canvas 或关闭 use_kg")
|
||||
|
||||
// ErrRetrievalServiceMissing is returned by the stub when the
|
||||
// internal/service/nlp RetrievalService is not wired. Plan §5 Phase 3
|
||||
// batch 1 ships the tool shell; service wiring lands in Phase 5.
|
||||
var ErrRetrievalServiceMissing = errors.New(
|
||||
"Retrieval service not yet implemented (Phase 5 wiring) — " +
|
||||
"use Python Canvas or implement internal/service/nlp/retrieval.go",
|
||||
)
|
||||
|
||||
// retrievalToolName preserves the Python typo ("dateset") for backward
|
||||
// compatibility with existing Canvas DSLs that reference the tool by name.
|
||||
const retrievalToolName = "search_my_dateset"
|
||||
|
||||
const retrievalToolDescription = "This tool can be utilized for relevant content searching in the datasets."
|
||||
|
||||
// retrievalArgs is the JSON schema the model sends into InvokableRun. We
|
||||
// accept both `query` (canonical) and `dataset_ids` / `use_kg` etc. to
|
||||
// match the Python ToolMeta field set.
|
||||
type retrievalArgs struct {
|
||||
Query string `json:"query"`
|
||||
DatasetIDs []string `json:"dataset_ids,omitempty"`
|
||||
TopN int `json:"top_n,omitempty"`
|
||||
UseKG bool `json:"use_kg,omitempty"`
|
||||
}
|
||||
|
||||
// retrievalResult is the JSON shape returned to the model. The `_ERROR`
|
||||
// field matches the Python tool's output convention; downstream components
|
||||
// can pattern-match on it.
|
||||
type retrievalResult struct {
|
||||
FormalizedContent string `json:"formalized_content,omitempty"`
|
||||
Chunks []chunkPayload `json:"chunks,omitempty"`
|
||||
Stub bool `json:"stub,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// chunkPayload is the minimal chunk shape we surface. We don't try to
|
||||
// match every Python field — the stub returns empty data; the wired
|
||||
// implementation will populate the real shape.
|
||||
type chunkPayload struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
DocumentID string `json:"document_id,omitempty"`
|
||||
Score float64 `json:"score,omitempty"`
|
||||
}
|
||||
|
||||
// RetrievalTool is the Phase 3 batch 1 shell for the Retrieval tool
|
||||
// (plan §2.11.4 row 15, §5 Phase 3 第 1 批).
|
||||
//
|
||||
// In Phase 3 batch 1 the tool is a STUB: it validates the input (rejecting
|
||||
// use_kg=true with ErrGraphRAGNotSupported) and returns a structured
|
||||
// "not-yet-wired" error so callers can detect the gap. Phase 5 wires the
|
||||
// in-process call to internal/service/nlp.RetrievalService.
|
||||
type RetrievalTool struct{}
|
||||
|
||||
// NewRetrievalTool returns a RetrievalTool implementing eino's
|
||||
// tool.InvokableTool interface.
|
||||
func NewRetrievalTool() *RetrievalTool {
|
||||
return &RetrievalTool{}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model. The schema mirrors
|
||||
// the Python RetrievalParam ToolMeta (plan §5 Phase 3, 字段对齐).
|
||||
func (r *RetrievalTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: retrievalToolName,
|
||||
Desc: retrievalToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "The keywords to search the dataset. The keywords should be the most important words/terms (including synonyms) from the original request.",
|
||||
Required: true,
|
||||
},
|
||||
"dataset_ids": {
|
||||
Type: schema.Array,
|
||||
Desc: "Optional list of dataset IDs to restrict the search to.",
|
||||
Required: false,
|
||||
},
|
||||
"top_n": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Number of top chunks to return. Defaults to 8 if omitted.",
|
||||
Required: false,
|
||||
},
|
||||
"use_kg": {
|
||||
Type: schema.Boolean,
|
||||
Desc: "GraphRAG toggle. Not supported in Go Canvas (plan §5 Phase 3); must be false.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// InvokableRun executes the tool. In Phase 3 batch 1 this validates the
|
||||
// input and returns a structured error (or ErrGraphRAGNotSupported) so
|
||||
// callers can detect the gap. Phase 5 will replace the stub body with
|
||||
// the in-process call to internal/service/nlp.RetrievalService.
|
||||
func (r *RetrievalTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
|
||||
var args retrievalArgs
|
||||
if argumentsInJSON != "" {
|
||||
if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil {
|
||||
return "", fmt.Errorf("retrieval: parse arguments: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if args.UseKG {
|
||||
// Plan §5 Phase 3 + §9 Q3: GraphRAG is out of scope for the Go
|
||||
// Canvas. Return the structured error so the model can react.
|
||||
return stubJSON(retrievalResult{
|
||||
Stub: true,
|
||||
Error: ErrGraphRAGNotSupported.Error(),
|
||||
}), ErrGraphRAGNotSupported
|
||||
}
|
||||
|
||||
// Phase 3 batch 1: in-process wiring lands in Phase 5. The stub
|
||||
// surfaces a clear, machine-detectable error.
|
||||
return stubJSON(retrievalResult{
|
||||
Stub: true,
|
||||
Error: ErrRetrievalServiceMissing.Error(),
|
||||
}), ErrRetrievalServiceMissing
|
||||
}
|
||||
|
||||
// stubJSON marshals the result and returns it as a string. Marshaling
|
||||
// failures are converted to a plain string error so the model can still
|
||||
// surface something to the user.
|
||||
func stubJSON(r retrievalResult) string {
|
||||
b, err := json.Marshal(r)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"retrieval: marshal stub result: %s","stub":true}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
108
internal/agent/tool/retrieval_test.go
Normal file
108
internal/agent/tool/retrieval_test.go
Normal file
@@ -0,0 +1,108 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rt := NewRetrievalTool()
|
||||
out, err := rt.InvokableRun(context.Background(), `{"query":"hello"}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected stub error, got nil")
|
||||
}
|
||||
if !errors.Is(err, ErrRetrievalServiceMissing) {
|
||||
t.Fatalf("err = %v, want ErrRetrievalServiceMissing", err)
|
||||
}
|
||||
|
||||
// Output is a JSON envelope with the stub error message.
|
||||
var got retrievalResult
|
||||
if jerr := json.Unmarshal([]byte(out), &got); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if !got.Stub {
|
||||
t.Errorf("Stub = false, want true")
|
||||
}
|
||||
if !strings.Contains(got.Error, "Phase 5 wiring") {
|
||||
t.Errorf("Error = %q, want to mention 'Phase 5 wiring'", got.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrieval_RejectsUseKG(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rt := NewRetrievalTool()
|
||||
out, err := rt.InvokableRun(context.Background(), `{"query":"x","use_kg":true}`)
|
||||
if !errors.Is(err, ErrGraphRAGNotSupported) {
|
||||
t.Fatalf("err = %v, want ErrGraphRAGNotSupported", err)
|
||||
}
|
||||
if !strings.Contains(out, "GraphRAG") {
|
||||
t.Errorf("output %q should mention GraphRAG", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrieval_InfoMatchesPythonMeta(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rt := NewRetrievalTool()
|
||||
info, err := rt.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "search_my_dateset" {
|
||||
t.Errorf("Name = %q, want search_my_dateset (typo preserved)", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "datasets") {
|
||||
t.Errorf("Desc = %q, want to mention 'datasets'", info.Desc)
|
||||
}
|
||||
// The query param must be present and required. ToJSONSchema returns
|
||||
// a *jsonschema.Schema whose Properties is an *orderedmap.Map; we use
|
||||
// MarshalJSON to assert the parameter set without depending on the
|
||||
// map's concrete Get signature.
|
||||
params, err := info.ToJSONSchema()
|
||||
if err != nil {
|
||||
t.Fatalf("ToJSONSchema: %v", err)
|
||||
}
|
||||
raw, err := json.Marshal(params)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal schema: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(raw), `"query"`) {
|
||||
t.Errorf("schema JSON does not contain 'query' key: %s", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetrieval_EmptyArgsIsHandled(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rt := NewRetrievalTool()
|
||||
// Empty arguments should still return a stub error (not panic) — the
|
||||
// Python tool defaults to empty_response in this case, but in Phase 3
|
||||
// batch 1 we don't have wiring, so we surface the service-missing
|
||||
// error.
|
||||
_, err := rt.InvokableRun(context.Background(), "")
|
||||
if !errors.Is(err, ErrRetrievalServiceMissing) {
|
||||
t.Fatalf("err = %v, want ErrRetrievalServiceMissing", err)
|
||||
}
|
||||
}
|
||||
168
internal/agent/tool/searxng.go
Normal file
168
internal/agent/tool/searxng.go
Normal file
@@ -0,0 +1,168 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const searxngToolName = "searxng"
|
||||
|
||||
const searxngToolDescription = "Search a self-hosted SearXNG instance. Returns results[].{title, url, content}."
|
||||
|
||||
// searxngParams is the JSON shape the model sends into InvokableRun.
|
||||
// base_url is the SearXNG root (no trailing slash); the default points
|
||||
// at a local instance on port 8888.
|
||||
type searxngParams struct {
|
||||
BaseURL string `json:"base_url"`
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
// searxngResult mirrors one element of the upstream `results` array.
|
||||
type searxngResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// searxngResponse is the upstream SearXNG JSON envelope.
|
||||
type searxngResponse struct {
|
||||
Results []searxngResult `json:"results"`
|
||||
}
|
||||
|
||||
// searxngEnvelope is the JSON shape the model sees.
|
||||
type searxngEnvelope struct {
|
||||
Results []searxngResult `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// SearXNGTool is the Phase 3 batch 2 implementation of the SearXNG
|
||||
// meta-search tool (plan §2.11.4 row 16, §5 Phase 3 第 2 批). It calls
|
||||
// a self-hosted SearXNG instance via the shared HTTPHelper.
|
||||
type SearXNGTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewSearXNGTool returns a SearXNGTool using the default HTTPHelper.
|
||||
func NewSearXNGTool() *SearXNGTool {
|
||||
return NewSearXNGToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewSearXNGToolWith returns a SearXNGTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewSearXNGToolWith(h *HTTPHelper) *SearXNGTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &SearXNGTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (s *SearXNGTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: searxngToolName,
|
||||
Desc: searxngToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query",
|
||||
Required: true,
|
||||
},
|
||||
"base_url": {
|
||||
Type: schema.String,
|
||||
Desc: `SearXNG base URL. Defaults to "http://localhost:8888".`,
|
||||
Required: false,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of results to return. Defaults to 5.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildSearXNGURL constructs the SearXNG /search URL. The base URL is
|
||||
// normalized to drop trailing slashes. We use url.JoinPath-style
|
||||
// construction (manual) so the function works on Go 1.19 as well.
|
||||
func buildSearXNGURL(baseURL, query string, maxResults int) string {
|
||||
if baseURL == "" {
|
||||
baseURL = "http://localhost:8888"
|
||||
}
|
||||
baseURL = strings.TrimRight(baseURL, "/")
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
q := url.Values{}
|
||||
q.Set("q", query)
|
||||
q.Set("format", "json")
|
||||
q.Set("language", "all")
|
||||
return baseURL + "/search?" + q.Encode()
|
||||
}
|
||||
|
||||
// InvokableRun performs the SearXNG search.
|
||||
func (s *SearXNGTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p searxngParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return searxngErrJSON(fmt.Errorf("searxng: parse arguments: %w", err)),
|
||||
fmt.Errorf("searxng: parse arguments: %w", err)
|
||||
}
|
||||
if p.Query == "" {
|
||||
return searxngErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("searxng: query is required")
|
||||
}
|
||||
|
||||
endpoint := buildSearXNGURL(p.BaseURL, p.Query, p.MaxResults)
|
||||
resp, err := s.helper.Do(ctx, http.MethodGet, endpoint, "", "", nil)
|
||||
if err != nil {
|
||||
return searxngErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return searxngErrJSON(fmt.Errorf("searxng: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("searxng: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw searxngResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return searxngErrJSON(fmt.Errorf("searxng: decode response: %w", err)),
|
||||
fmt.Errorf("searxng: decode response: %w", err)
|
||||
}
|
||||
return searxngJSON(searxngEnvelope{Results: raw.Results}), nil
|
||||
}
|
||||
|
||||
func searxngJSON(env searxngEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"searxng: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func searxngErrJSON(err error) string {
|
||||
return searxngJSON(searxngEnvelope{Error: err.Error()})
|
||||
}
|
||||
156
internal/agent/tool/searxng_test.go
Normal file
156
internal/agent/tool/searxng_test.go
Normal file
@@ -0,0 +1,156 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSearXNG_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
base string
|
||||
query string
|
||||
max int
|
||||
wantPath string
|
||||
wantHost string
|
||||
}{
|
||||
{
|
||||
name: "default base",
|
||||
base: "",
|
||||
query: "ragflow",
|
||||
max: 0,
|
||||
wantPath: "/search",
|
||||
wantHost: "localhost:8888",
|
||||
},
|
||||
{
|
||||
name: "custom base trailing slash",
|
||||
base: "https://searx.example.com/",
|
||||
query: "rag",
|
||||
max: 3,
|
||||
wantPath: "/search",
|
||||
wantHost: "searx.example.com",
|
||||
},
|
||||
{
|
||||
name: "custom base no slash",
|
||||
base: "https://searx.example.com",
|
||||
query: "rag",
|
||||
max: 7,
|
||||
wantPath: "/search",
|
||||
wantHost: "searx.example.com",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildSearXNGURL(tc.base, tc.query, tc.max)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != tc.wantHost {
|
||||
t.Errorf("host = %q, want %q", u.Host, tc.wantHost)
|
||||
}
|
||||
if u.Path != tc.wantPath {
|
||||
t.Errorf("path = %q, want %q", u.Path, tc.wantPath)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("q") != tc.query {
|
||||
t.Errorf("q = %q, want %q", q.Get("q"), tc.query)
|
||||
}
|
||||
if q.Get("format") != "json" {
|
||||
t.Errorf("format = %q, want json", q.Get("format"))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNG_ParseResults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"results": [
|
||||
{"title":"RAGFlow","url":"https://ragflow.io","content":"Open source RAG engine"},
|
||||
{"title":"GitHub","url":"https://github.com/infiniflow/ragflow","content":"Source code"}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
tool := NewSearXNGTool()
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"ragflow","base_url":`+jsonString(srv.URL)+`,"max_results":5}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env searxngEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
}
|
||||
if env.Results[0].Title != "RAGFlow" {
|
||||
t.Errorf("Results[0].Title = %q, want RAGFlow", env.Results[0].Title)
|
||||
}
|
||||
if env.Results[1].URL != "https://github.com/infiniflow/ragflow" {
|
||||
t.Errorf("Results[1].URL = %q, want https://github.com/infiniflow/ragflow", env.Results[1].URL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNG_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewSearXNGTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "searxng" {
|
||||
t.Errorf("Name = %q, want searxng", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "SearXNG") {
|
||||
t.Errorf("Desc = %q, want to mention SearXNG", info.Desc)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNG_RequiresQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewSearXNGTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{"query":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty query")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "query") {
|
||||
t.Errorf("err = %v, want to mention query", err)
|
||||
}
|
||||
}
|
||||
162
internal/agent/tool/ssrf.go
Normal file
162
internal/agent/tool/ssrf.go
Normal file
@@ -0,0 +1,162 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ErrSSRFBlocked is returned when a tool is asked to fetch a URL whose
|
||||
// host resolves to a loopback, private, link-local or otherwise
|
||||
// non-public IP range. This blocks the standard SSRF probes against
|
||||
// internal services and cloud metadata endpoints (AWS, GCP, Azure,
|
||||
// Alibaba all expose 169.254.169.254).
|
||||
var ErrSSRFBlocked = errors.New("ssrf: target host is blocked by SSRF guard")
|
||||
|
||||
// credentialQueryParams is the lower-cased set of query-parameter names
|
||||
// we treat as API credentials. Matching is case-insensitive. Any query
|
||||
// string that uses one of these names has its value redacted in error
|
||||
// messages and logs so an upstream 4xx/5xx that echoes the URL never
|
||||
// leaks the secret.
|
||||
var credentialQueryParams = map[string]struct{}{
|
||||
"key": {},
|
||||
"api_key": {},
|
||||
"apikey": {},
|
||||
"token": {},
|
||||
"access_token": {},
|
||||
"auth": {},
|
||||
}
|
||||
|
||||
// validateURLForSSRF parses rawURL and rejects any target whose host
|
||||
// resolves (via DNS) to a non-public IP. The check is repeated against
|
||||
// every returned A/AAAA record because a name that resolves to a mix
|
||||
// of public and private IPs is also dangerous (DNS-rebinding /
|
||||
// multi-A-record pinning).
|
||||
func validateURLForSSRF(rawURL string) error {
|
||||
_, _, err := ResolveAndValidate(rawURL)
|
||||
return err
|
||||
}
|
||||
|
||||
// ResolveAndValidate parses rawURL, performs the SSRF blocklist checks,
|
||||
// and returns the first non-public IP that the host resolves to. The
|
||||
// returned IP is safe to dial directly (bypassing a fresh DNS lookup)
|
||||
// which defeats DNS-rebinding attacks: an attacker cannot swap a
|
||||
// public record for a private one between this lookup and the connect,
|
||||
// because the connect is pinned at the transport layer (see
|
||||
// HTTPHelper.DoPinned in http_helper.go) and never re-resolves the
|
||||
// hostname. Callers feed (originalHost, pinnedIP) into DoPinned.
|
||||
//
|
||||
// Note: pinning is done at the *http.Transport dialer, not by mutating
|
||||
// the request URL. Mutating u.Host to the IP would break HTTPS — TLS
|
||||
// ServerName is auto-populated from req.URL.Host, so the SNI would
|
||||
// become the IP and cert verification would target the IP. The
|
||||
// transport-layer approach keeps the URL host as the original hostname,
|
||||
// preserving correct SNI / cert verification for any HTTPS endpoint.
|
||||
func ResolveAndValidate(rawURL string) (originalHost string, pinnedIP net.IP, err error) {
|
||||
u, perr := url.Parse(rawURL)
|
||||
if perr != nil {
|
||||
return "", nil, fmt.Errorf("ssrf: parse url: %w", perr)
|
||||
}
|
||||
if u.Scheme != "http" && u.Scheme != "https" {
|
||||
return "", nil, fmt.Errorf("ssrf: unsupported scheme %q", u.Scheme)
|
||||
}
|
||||
host := u.Hostname()
|
||||
if host == "" {
|
||||
return "", nil, fmt.Errorf("ssrf: empty host")
|
||||
}
|
||||
|
||||
// Short-circuit the well-known host aliases that DNS lookups may
|
||||
// also catch, but defending against the literal name is cheap and
|
||||
// saves a syscall on the common probe path.
|
||||
lower := strings.ToLower(host)
|
||||
if lower == "localhost" || strings.HasSuffix(lower, ".localhost") ||
|
||||
lower == "metadata.google.internal" || lower == "metadata" ||
|
||||
lower == "0.0.0.0" || lower == "::" {
|
||||
return "", nil, fmt.Errorf("%w: %s", ErrSSRFBlocked, host)
|
||||
}
|
||||
|
||||
// If the host is a literal IP, no DNS lookup is needed.
|
||||
if ip := net.ParseIP(host); ip != nil {
|
||||
if isPrivateOrLoopback(ip) {
|
||||
return "", nil, fmt.Errorf("%w: literal %s", ErrSSRFBlocked, host)
|
||||
}
|
||||
return host, ip, nil
|
||||
}
|
||||
|
||||
ips, lerr := net.LookupIP(host)
|
||||
if lerr != nil {
|
||||
return "", nil, fmt.Errorf("ssrf: resolve %s: %w", host, lerr)
|
||||
}
|
||||
var firstSafe net.IP
|
||||
for _, ip := range ips {
|
||||
if isPrivateOrLoopback(ip) {
|
||||
return "", nil, fmt.Errorf("%w: %s -> %s", ErrSSRFBlocked, host, ip)
|
||||
}
|
||||
if firstSafe == nil {
|
||||
firstSafe = ip
|
||||
}
|
||||
}
|
||||
if firstSafe == nil {
|
||||
return "", nil, fmt.Errorf("ssrf: %s has no A/AAAA records", host)
|
||||
}
|
||||
return host, firstSafe, nil
|
||||
}
|
||||
|
||||
// isPrivateOrLoopback reports whether ip is in any of the ranges we
|
||||
// refuse to fetch from. It is deliberately conservative — link-local
|
||||
// (169.254.0.0/16, fe80::/10) is rejected because that is where cloud
|
||||
// metadata services live; multicast and the unspecified address are
|
||||
// also rejected.
|
||||
func isPrivateOrLoopback(ip net.IP) bool {
|
||||
if ip == nil {
|
||||
return true
|
||||
}
|
||||
if ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() ||
|
||||
ip.IsPrivate() || ip.IsMulticast() || ip.IsUnspecified() {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// SanitizeURL strips query parameters whose names match a small set of
|
||||
// well-known credential names so error messages and logs that echo the
|
||||
// request URL do not leak API keys. Anything else is preserved. The
|
||||
// returned string is always a valid URL; on parse failure the original
|
||||
// is returned unchanged.
|
||||
func SanitizeURL(rawURL string) string {
|
||||
u, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return rawURL
|
||||
}
|
||||
q := u.Query()
|
||||
changed := false
|
||||
for k := range q {
|
||||
if _, ok := credentialQueryParams[strings.ToLower(k)]; ok {
|
||||
q.Set(k, "REDACTED")
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return rawURL
|
||||
}
|
||||
u.RawQuery = q.Encode()
|
||||
return u.String()
|
||||
}
|
||||
241
internal/agent/tool/ssrf_test.go
Normal file
241
internal/agent/tool/ssrf_test.go
Normal file
@@ -0,0 +1,241 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"net"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestValidateURLForSSRF(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
rawURL string
|
||||
wantErr bool
|
||||
}{
|
||||
// Public targets: should pass (DNS may still fail in CI; we
|
||||
// accept both "ok" and "resolve error").
|
||||
{"https_public", "https://example.com/", false},
|
||||
{"http_public", "http://1.1.1.1/", false},
|
||||
|
||||
// Loopback / metadata / private: must be blocked.
|
||||
{"localhost_alias", "http://localhost/", true},
|
||||
{"localhost_subdomain", "http://admin.localhost/", true},
|
||||
{"loopback_ip_v4", "http://127.0.0.1/", true},
|
||||
{"loopback_ip_v6", "http://[::1]/", true},
|
||||
{"unspecified_v4", "http://0.0.0.0/", true},
|
||||
{"unspecified_v6", "http://[::]/", true},
|
||||
{"private_10", "http://10.0.0.1/", true},
|
||||
{"private_192", "http://192.168.1.1/", true},
|
||||
{"private_172", "http://172.16.0.1/", true},
|
||||
{"link_local_v4_metadata", "http://169.254.169.254/latest/meta-data/", true},
|
||||
{"link_local_v6", "http://[fe80::1]/", true},
|
||||
{"multicast", "http://224.0.0.1/", true},
|
||||
{"gcp_metadata_alias", "http://metadata.google.internal/", true},
|
||||
{"metadata_short", "http://metadata/", true},
|
||||
|
||||
// Bad input.
|
||||
{"empty", "", true},
|
||||
{"no_scheme", "example.com/foo", true},
|
||||
{"file_scheme", "file:///etc/passwd", true},
|
||||
{"javascript_scheme", "javascript:alert(1)", true},
|
||||
{"no_host", "http:///path", true},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := validateURLForSSRF(tc.rawURL)
|
||||
if tc.wantErr {
|
||||
if err == nil {
|
||||
t.Fatalf("validateURLForSSRF(%q) = nil, want error", tc.rawURL)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "ssrf") && tc.rawURL != "" && tc.rawURL != "file:///etc/passwd" {
|
||||
t.Errorf("validateURLForSSRF(%q) = %q, want an ssrf-tagged error", tc.rawURL, err)
|
||||
}
|
||||
} else if err != nil {
|
||||
// Public target may fail DNS in CI; tolerate resolve errors.
|
||||
if !strings.Contains(err.Error(), "resolve") {
|
||||
t.Errorf("validateURLForSSRF(%q) = %q, want nil or resolve error", tc.rawURL, err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveAndValidate covers the resolver half of the M1-rebinding
|
||||
// fix: for a literal public IP the function returns that IP unchanged;
|
||||
// for a hostname it returns the first A/AAAA record (which is what
|
||||
// DoPinned will then dial directly, defeating DNS rebinding between
|
||||
// validation and connect). For all SSRF-blocked targets the function
|
||||
// must return an error AND a nil IP, so DoPinned cannot be invoked
|
||||
// with a stale allow-list entry.
|
||||
func TestResolveAndValidate(t *testing.T) {
|
||||
t.Run("literal_public_ip", func(t *testing.T) {
|
||||
host, ip, err := ResolveAndValidate("https://1.1.1.1/foo")
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveAndValidate(literal public IP) = %v, want nil", err)
|
||||
}
|
||||
if host != "1.1.1.1" {
|
||||
t.Errorf("host = %q, want 1.1.1.1", host)
|
||||
}
|
||||
if ip == nil || !ip.Equal(net.ParseIP("1.1.1.1")) {
|
||||
t.Errorf("ip = %v, want 1.1.1.1", ip)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("literal_blocked_ip", func(t *testing.T) {
|
||||
host, ip, err := ResolveAndValidate("http://127.0.0.1/")
|
||||
if err == nil {
|
||||
t.Fatalf("ResolveAndValidate(loopback) = nil, want error")
|
||||
}
|
||||
if host != "" {
|
||||
t.Errorf("host = %q, want empty on error", host)
|
||||
}
|
||||
if ip != nil {
|
||||
t.Errorf("ip = %v, want nil on error", ip)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("literal_blocked_metadata", func(t *testing.T) {
|
||||
_, ip, err := ResolveAndValidate("http://169.254.169.254/latest/meta-data/")
|
||||
if err == nil {
|
||||
t.Fatalf("ResolveAndValidate(metadata) = nil, want error")
|
||||
}
|
||||
if ip != nil {
|
||||
t.Errorf("ip = %v, want nil on error", ip)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("hostname_resolves_to_public", func(t *testing.T) {
|
||||
// example.com is required to resolve to a public IP per RFC 2606.
|
||||
// The DNS lookup is wrapped in a goroutine with a 2s deadline so
|
||||
// sandboxed CI environments without upstream DNS can skip the
|
||||
// test rather than hang the suite (the default LookupIP honours
|
||||
// no timeout; this was the root cause of the 60s test timeout
|
||||
// during the rebinding-hardening review).
|
||||
type result struct {
|
||||
host string
|
||||
ip net.IP
|
||||
err error
|
||||
}
|
||||
ch := make(chan result, 1)
|
||||
go func() {
|
||||
h, ip, err := ResolveAndValidate("https://example.com/")
|
||||
ch <- result{h, ip, err}
|
||||
}()
|
||||
select {
|
||||
case r := <-ch:
|
||||
if r.err != nil {
|
||||
if strings.Contains(r.err.Error(), "resolve") {
|
||||
t.Skipf("DNS unavailable in CI: %v", r.err)
|
||||
}
|
||||
t.Fatalf("ResolveAndValidate(example.com) = %v, want nil", r.err)
|
||||
}
|
||||
if r.host != "example.com" {
|
||||
t.Errorf("host = %q, want example.com", r.host)
|
||||
}
|
||||
if r.ip == nil {
|
||||
t.Fatalf("ip = nil, want non-nil")
|
||||
}
|
||||
if r.ip.IsLoopback() || r.ip.IsPrivate() || r.ip.IsLinkLocalUnicast() {
|
||||
t.Errorf("ip = %s, want a public address", r.ip)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Skip("DNS lookup for example.com timed out — sandboxed CI without upstream DNS")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("hostname_blocked_alias", func(t *testing.T) {
|
||||
_, ip, err := ResolveAndValidate("http://localhost/")
|
||||
if err == nil {
|
||||
t.Fatalf("ResolveAndValidate(localhost) = nil, want error")
|
||||
}
|
||||
if ip != nil {
|
||||
t.Errorf("ip = %v, want nil on error", ip)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("bad_input_returns_nil_ip", func(t *testing.T) {
|
||||
for _, bad := range []string{"", "file:///etc/passwd", "http:///path", "example.com/no-scheme"} {
|
||||
_, ip, err := ResolveAndValidate(bad)
|
||||
if err == nil {
|
||||
t.Errorf("ResolveAndValidate(%q) = nil err, want error", bad)
|
||||
}
|
||||
if ip != nil {
|
||||
t.Errorf("ResolveAndValidate(%q) ip = %v, want nil on error", bad, ip)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSanitizeURL(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
raw string
|
||||
want string
|
||||
noKeys bool
|
||||
}{
|
||||
{
|
||||
name: "google_api_key_in_query",
|
||||
raw: "https://www.googleapis.com/customsearch/v1?key=AIzaSecret&q=hello",
|
||||
want: "https://www.googleapis.com/customsearch/v1?key=REDACTED&q=hello",
|
||||
},
|
||||
{
|
||||
name: "qweather_key",
|
||||
raw: "https://devapi.qweather.com/v7/weather/now?location=101010100&key=abc123&lang=zh",
|
||||
want: "https://devapi.qweather.com/v7/weather/now?key=REDACTED&lang=zh&location=101010100",
|
||||
},
|
||||
{
|
||||
name: "token_param",
|
||||
raw: "https://api.example.com/data?token=topsecret&page=2",
|
||||
want: "https://api.example.com/data?page=2&token=REDACTED",
|
||||
},
|
||||
{
|
||||
name: "no_creds_unchanged",
|
||||
raw: "https://api.example.com/data?page=2&size=10",
|
||||
want: "https://api.example.com/data?page=2&size=10",
|
||||
noKeys: true,
|
||||
},
|
||||
{
|
||||
name: "unparseable_unchanged",
|
||||
raw: "://broken",
|
||||
want: "://broken",
|
||||
noKeys: true,
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := SanitizeURL(tc.raw)
|
||||
if tc.noKeys {
|
||||
if got != tc.raw {
|
||||
t.Errorf("SanitizeURL(%q) = %q, want unchanged", tc.raw, got)
|
||||
}
|
||||
return
|
||||
}
|
||||
// Compare on a normalised form by stripping the input URL
|
||||
// and re-encoding both — easier than asserting exact
|
||||
// query-string order, which is non-deterministic in net/url.
|
||||
if !strings.Contains(got, "REDACTED") {
|
||||
t.Errorf("SanitizeURL(%q) = %q, want REDACTED marker", tc.raw, got)
|
||||
}
|
||||
if strings.Contains(got, "AIzaSecret") || strings.Contains(got, "abc123") || strings.Contains(got, "topsecret") {
|
||||
t.Errorf("SanitizeURL(%q) = %q, secret value still present", tc.raw, got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
215
internal/agent/tool/tavily.go
Normal file
215
internal/agent/tool/tavily.go
Normal file
@@ -0,0 +1,215 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const tavilyToolName = "tavily"
|
||||
|
||||
const tavilyToolDescription = "Search the web via the Tavily API. Returns a list of {url, title, content} results."
|
||||
|
||||
// tavilyParams is the JSON shape the model sends into InvokableRun. The
|
||||
// api_key may be omitted when the env var TAVILY_API_KEY is set; the tool
|
||||
// resolves it from the environment in that case.
|
||||
type tavilyParams struct {
|
||||
APIKey string `json:"api_key"`
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
SearchDepth string `json:"search_depth"`
|
||||
}
|
||||
|
||||
// tavilyRequestBody is the JSON body POSTed to the Tavily /search endpoint.
|
||||
// The struct matches the upstream API (https://docs.tavily.com).
|
||||
type tavilyRequestBody struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
SearchDepth string `json:"search_depth"`
|
||||
}
|
||||
|
||||
// tavilyResult mirrors one element of the upstream `results` array. We
|
||||
// return these verbatim to the model.
|
||||
type tavilyResult struct {
|
||||
URL string `json:"url"`
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// tavilyResponse is the envelope returned by Tavily. We only model the
|
||||
// fields we care about; the upstream API has more, but they are ignored.
|
||||
type tavilyResponse struct {
|
||||
Results []tavilyResult `json:"results"`
|
||||
}
|
||||
|
||||
// tavilyEnvelope is the shape the model actually sees, identical to the
|
||||
// Python tool's output convention.
|
||||
type tavilyEnvelope struct {
|
||||
Results []tavilyResult `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// TavilyTool is the Phase 3 batch 2 implementation of the Tavily search
|
||||
// tool (plan §2.11.4 row 17, §5 Phase 3 第 2 批). It POSTs a search request
|
||||
// to https://api.tavily.com/search using the shared HTTPHelper and returns
|
||||
// the upstream `results` array as JSON.
|
||||
type TavilyTool struct {
|
||||
helper *HTTPHelper
|
||||
envKey func() string
|
||||
}
|
||||
|
||||
// NewTavilyTool returns a TavilyTool using the default HTTPHelper and
|
||||
// the TAVILY_API_KEY env var for credential resolution.
|
||||
func NewTavilyTool() *TavilyTool {
|
||||
return NewTavilyToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewTavilyToolWith returns a TavilyTool that uses the provided
|
||||
// HTTPHelper. Useful for tests that want to inject a custom transport.
|
||||
func NewTavilyToolWith(h *HTTPHelper) *TavilyTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &TavilyTool{helper: h, envKey: defaultTavilyEnvKey}
|
||||
}
|
||||
|
||||
// NewTavilyToolWithEnvKey returns a TavilyTool with a custom env-key
|
||||
// resolver. Useful for tests that want to inject a fake credential
|
||||
// without mutating process state.
|
||||
func NewTavilyToolWithEnvKey(h *HTTPHelper, envKey func() string) *TavilyTool {
|
||||
if envKey == nil {
|
||||
envKey = defaultTavilyEnvKey
|
||||
}
|
||||
return &TavilyTool{helper: h, envKey: envKey}
|
||||
}
|
||||
|
||||
// defaultTavilyEnvKey is the production env-key resolver. Pulled out
|
||||
// as a named function (not a var) so tests cannot accidentally
|
||||
// mutate it via package-var assignment.
|
||||
func defaultTavilyEnvKey() string { return os.Getenv("TAVILY_API_KEY") }
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (t *TavilyTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: tavilyToolName,
|
||||
Desc: tavilyToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query",
|
||||
Required: true,
|
||||
},
|
||||
"api_key": {
|
||||
Type: schema.String,
|
||||
Desc: "Tavily API key. Falls back to TAVILY_API_KEY env var.",
|
||||
Required: false,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of results to return. Defaults to 5.",
|
||||
Required: false,
|
||||
},
|
||||
"search_depth": {
|
||||
Type: schema.String,
|
||||
Desc: `Tavily search depth: "basic" (default) or "advanced".`,
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// tavilyEndpoint is the Tavily /search URL. Exposed as a package var so
|
||||
// tests can substitute a httptest.Server URL. Tests must serialize
|
||||
// access with tavilyEndpointMu if running in parallel.
|
||||
var tavilyEndpoint = "https://api.tavily.com/search"
|
||||
|
||||
// InvokableRun performs the Tavily search. The api_key may come from the
|
||||
// argument or the TAVILY_API_KEY env var.
|
||||
func (t *TavilyTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p tavilyParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return tavilyErrJSON(fmt.Errorf("tavily: parse arguments: %w", err)),
|
||||
fmt.Errorf("tavily: parse arguments: %w", err)
|
||||
}
|
||||
if p.Query == "" {
|
||||
return tavilyErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("tavily: query is required")
|
||||
}
|
||||
if p.MaxResults <= 0 {
|
||||
p.MaxResults = 5
|
||||
}
|
||||
if p.SearchDepth == "" {
|
||||
p.SearchDepth = "basic"
|
||||
}
|
||||
|
||||
apiKey := p.APIKey
|
||||
if apiKey == "" {
|
||||
apiKey = t.envKey()
|
||||
}
|
||||
if apiKey == "" {
|
||||
return tavilyErrJSON(fmt.Errorf("tavily: api_key is required (or set TAVILY_API_KEY)")),
|
||||
fmt.Errorf("tavily: api_key is required (or set TAVILY_API_KEY)")
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(tavilyRequestBody{
|
||||
Query: p.Query,
|
||||
MaxResults: p.MaxResults,
|
||||
SearchDepth: p.SearchDepth,
|
||||
})
|
||||
|
||||
resp, err := t.helper.Do(ctx,
|
||||
http.MethodPost, tavilyEndpoint, string(body), "application/json",
|
||||
map[string]string{"Authorization": "Bearer " + apiKey},
|
||||
)
|
||||
if err != nil {
|
||||
return tavilyErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return tavilyErrJSON(fmt.Errorf("tavily: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("tavily: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw tavilyResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return tavilyErrJSON(fmt.Errorf("tavily: decode response: %w", err)),
|
||||
fmt.Errorf("tavily: decode response: %w", err)
|
||||
}
|
||||
return tavilyJSON(tavilyEnvelope{Results: raw.Results}), nil
|
||||
}
|
||||
|
||||
// tavilyJSON marshals the envelope to a JSON string for the model.
|
||||
func tavilyJSON(env tavilyEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"tavily: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// tavilyErrJSON wraps an error in the standard envelope.
|
||||
func tavilyErrJSON(err error) string {
|
||||
return tavilyJSON(tavilyEnvelope{Error: err.Error()})
|
||||
}
|
||||
178
internal/agent/tool/tavily_test.go
Normal file
178
internal/agent/tool/tavily_test.go
Normal file
@@ -0,0 +1,178 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTavily_BuildRequest(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var gotPath, gotAuth, gotCT, gotMethod string
|
||||
var gotBody map[string]any
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
gotMethod = r.Method
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
gotCT = r.Header.Get("Content-Type")
|
||||
_ = json.NewDecoder(r.Body).Decode(&gotBody)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"results":[]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
// Point the hard-coded tavily endpoint at the test server via a
|
||||
// transport that rewrites the host. Avoids the global-package-var
|
||||
// race that breaks parallel tests.
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewTavilyToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"ragflow","api_key":"key-xyz","max_results":3,"search_depth":"advanced"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if out == "" {
|
||||
t.Fatal("InvokableRun returned empty string")
|
||||
}
|
||||
|
||||
if gotMethod != http.MethodPost {
|
||||
t.Errorf("method = %q, want POST", gotMethod)
|
||||
}
|
||||
if !strings.HasSuffix(gotPath, "search") {
|
||||
t.Errorf("path = %q, want .../search", gotPath)
|
||||
}
|
||||
if gotAuth != "Bearer key-xyz" {
|
||||
t.Errorf("Authorization = %q, want %q", gotAuth, "Bearer key-xyz")
|
||||
}
|
||||
if !strings.HasPrefix(gotCT, "application/json") {
|
||||
t.Errorf("Content-Type = %q, want application/json", gotCT)
|
||||
}
|
||||
if gotBody["query"] != "ragflow" {
|
||||
t.Errorf("body.query = %v, want ragflow", gotBody["query"])
|
||||
}
|
||||
if n, ok := gotBody["max_results"].(float64); !ok || int(n) != 3 {
|
||||
t.Errorf("body.max_results = %v, want 3", gotBody["max_results"])
|
||||
}
|
||||
if gotBody["search_depth"] != "advanced" {
|
||||
t.Errorf("body.search_depth = %v, want advanced", gotBody["search_depth"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestTavily_ParseResponse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"results": [
|
||||
{"url":"https://a.example/","title":"A","content":"alpha"},
|
||||
{"url":"https://b.example/","title":"B","content":"beta"}
|
||||
],
|
||||
"answer": "ignored"
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewTavilyToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"x","api_key":"k"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env tavilyEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
}
|
||||
if env.Results[0].URL != "https://a.example/" || env.Results[0].Title != "A" {
|
||||
t.Errorf("Results[0] = %+v, want url=https://a.example/ title=A", env.Results[0])
|
||||
}
|
||||
if env.Results[1].Content != "beta" {
|
||||
t.Errorf("Results[1].Content = %q, want beta", env.Results[1].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTavily_RequiresAPIKey(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// envKey always returns "" so we know the failure is from the
|
||||
// missing api_key, not from a stray process env var.
|
||||
tool := NewTavilyToolWithEnvKey(NewHTTPHelper(), func() string { return "" })
|
||||
_, err := tool.InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for missing api_key")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "api_key") {
|
||||
t.Errorf("err = %v, want to mention api_key", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTavily_APIKeyFromEnv(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var gotAuth string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
_, _ = w.Write([]byte(`{"results":[]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewTavilyToolWithEnvKey(helper, func() string { return "from-env" })
|
||||
if _, err := tool.InvokableRun(context.Background(), `{"query":"x"}`); err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if gotAuth != "Bearer from-env" {
|
||||
t.Errorf("Authorization = %q, want Bearer from-env", gotAuth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTavily_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewTavilyTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "tavily" {
|
||||
t.Errorf("Name = %q, want tavily", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Tavily") {
|
||||
t.Errorf("Desc = %q, want to mention Tavily", info.Desc)
|
||||
}
|
||||
}
|
||||
251
internal/agent/tool/tushare.go
Normal file
251
internal/agent/tool/tushare.go
Normal file
@@ -0,0 +1,251 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const tushareToolName = "tushare"
|
||||
|
||||
const tushareToolDescription = "Call Tushare Pro (api.tushare.pro) for Chinese financial data. " +
|
||||
"Returns the upstream data.fields and data.items arrays."
|
||||
|
||||
// tushareEndpoint is the Tushare Pro REST endpoint. Exposed as a
|
||||
// package var so tests can substitute a httptest.Server URL.
|
||||
var tushareEndpoint = "http://api.tushare.pro"
|
||||
|
||||
// tushareParams is the JSON shape the model sends into InvokableRun.
|
||||
//
|
||||
// - Token (required): the Tushare Pro API token (积分账户).
|
||||
// - APIName (required): the Tushare interface name, e.g. "stock_basic",
|
||||
// "daily", "fund_basic", "index_daily", "fut_basic".
|
||||
// - Params (optional): a free-form map of Tushare query parameters,
|
||||
// e.g. {"ts_code":"000001.SZ","start_date":"20240101"}.
|
||||
type tushareParams struct {
|
||||
Token string `json:"token"`
|
||||
APIName string `json:"api_name"`
|
||||
Params map[string]string `json:"params,omitempty"`
|
||||
Fields []string `json:"fields,omitempty"`
|
||||
}
|
||||
|
||||
// tushareRequest is the JSON envelope Tushare Pro expects on POST.
|
||||
type tushareRequest struct {
|
||||
Token string `json:"token"`
|
||||
APIName string `json:"api_name"`
|
||||
Params map[string]string `json:"params,omitempty"`
|
||||
Fields []string `json:"fields,omitempty"`
|
||||
}
|
||||
|
||||
// tushareData is the upstream `data` field shape. Tushare returns a
|
||||
// column-major record: `fields` lists the column names in order, and
|
||||
// `items` is a slice of rows where each row is a slice aligned with
|
||||
// `fields`.
|
||||
type tushareData struct {
|
||||
Fields []string `json:"fields"`
|
||||
Items [][]any `json:"items"`
|
||||
}
|
||||
|
||||
// tushareResponse is the upstream Tushare Pro envelope.
|
||||
//
|
||||
// {
|
||||
// "code": 0, // 0 = OK
|
||||
// "msg": "...",
|
||||
// "data": {...} // optional
|
||||
// }
|
||||
type tushareResponse struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg,omitempty"`
|
||||
Data *tushareData `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
// tushareEnvelope is what the model sees. We pass through fields/items
|
||||
// verbatim so the model can index by column name. _ERROR captures
|
||||
// non-zero `code` responses from Tushare (e.g. "权限不足" / 40201) and
|
||||
// transport-level failures.
|
||||
type tushareEnvelope struct {
|
||||
Fields []string `json:"fields,omitempty"`
|
||||
Items [][]any `json:"items,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// TushareTool is the Phase 3 batch 4 implementation of the Tushare Pro
|
||||
// Chinese financial data tool (plan §2.11.4 row 19 [numbered 19, also
|
||||
// mentioned as row 18 in the matrix header], §5 Phase 3 第 4 批). It
|
||||
// performs a POST against api.tushare.pro with the token in the body
|
||||
// and returns the data.fields + data.items arrays.
|
||||
//
|
||||
// TushareTool uses the shared HTTPHelper for retry/timeout/OTel
|
||||
// propagation.
|
||||
type TushareTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewTushareTool returns a TushareTool using the default HTTPHelper.
|
||||
func NewTushareTool() *TushareTool {
|
||||
return NewTushareToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewTushareToolWith returns a TushareTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewTushareToolWith(h *HTTPHelper) *TushareTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &TushareTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (t *TushareTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: tushareToolName,
|
||||
Desc: tushareToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"token": {
|
||||
Type: schema.String,
|
||||
Desc: "Tushare Pro API token (积分账户).",
|
||||
Required: true,
|
||||
},
|
||||
"api_name": {
|
||||
Type: schema.String,
|
||||
Desc: "Tushare interface name, e.g. stock_basic, daily, fund_basic, index_daily, fut_basic.",
|
||||
Required: true,
|
||||
},
|
||||
"params": {
|
||||
Type: schema.Object,
|
||||
Desc: "Optional query parameters, e.g. {\"ts_code\":\"000001.SZ\",\"start_date\":\"20240101\"}.",
|
||||
Required: false,
|
||||
},
|
||||
"fields": {
|
||||
Type: schema.Array,
|
||||
Desc: "Optional subset of columns to return (Tushare `fields` parameter).",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildTushareRequestBody marshals the request envelope to JSON. The
|
||||
// Tushare Pro server expects a flat object — no URL-encoding — so we
|
||||
// POST raw JSON. Exposed for unit testing.
|
||||
func buildTushareRequestBody(p tushareParams) ([]byte, error) {
|
||||
req := tushareRequest{
|
||||
Token: p.Token,
|
||||
APIName: p.APIName,
|
||||
Params: p.Params,
|
||||
Fields: p.Fields,
|
||||
}
|
||||
return json.Marshal(req)
|
||||
}
|
||||
|
||||
// buildTushareURL returns the POST URL. Tushare's API only reads the
|
||||
// body, not the query string, but the helper requires a URL.
|
||||
func buildTushareURL() string {
|
||||
u, _ := url.Parse(tushareEndpoint)
|
||||
if u.Scheme == "" {
|
||||
u.Scheme = "http"
|
||||
}
|
||||
return u.String()
|
||||
}
|
||||
|
||||
// InvokableRun performs the Tushare Pro POST call.
|
||||
func (t *TushareTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p tushareParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return tushareErrJSON(fmt.Errorf("tushare: parse arguments: %w", err)),
|
||||
fmt.Errorf("tushare: parse arguments: %w", err)
|
||||
}
|
||||
if p.Token == "" {
|
||||
return tushareErrJSON(fmt.Errorf("tushare: token is required")),
|
||||
fmt.Errorf("tushare: token is required")
|
||||
}
|
||||
if p.APIName == "" {
|
||||
return tushareErrJSON(fmt.Errorf("tushare: api_name is required")),
|
||||
fmt.Errorf("tushare: api_name is required")
|
||||
}
|
||||
|
||||
body, err := buildTushareRequestBody(p)
|
||||
if err != nil {
|
||||
return tushareErrJSON(fmt.Errorf("tushare: build request: %w", err)),
|
||||
fmt.Errorf("tushare: build request: %w", err)
|
||||
}
|
||||
|
||||
headers := map[string]string{
|
||||
"Accept": "application/json",
|
||||
}
|
||||
|
||||
resp, err := t.helper.Do(ctx, http.MethodPost, buildTushareURL(), string(body), "application/json", headers)
|
||||
if err != nil {
|
||||
return tushareErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return tushareErrJSON(fmt.Errorf("tushare: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("tushare: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw tushareResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return tushareErrJSON(fmt.Errorf("tushare: decode response: %w", err)),
|
||||
fmt.Errorf("tushare: decode response: %w", err)
|
||||
}
|
||||
if raw.Code != 0 {
|
||||
msg := strings.TrimSpace(raw.Msg)
|
||||
if msg == "" {
|
||||
msg = fmt.Sprintf("tushare: upstream returned code %d", raw.Code)
|
||||
}
|
||||
return tushareErrJSON(fmt.Errorf("tushare: %s", msg)),
|
||||
fmt.Errorf("tushare: %s", msg)
|
||||
}
|
||||
|
||||
env := tushareEnvelope{Error: ""}
|
||||
if raw.Data != nil {
|
||||
env.Fields = raw.Data.Fields
|
||||
env.Items = raw.Data.Items
|
||||
}
|
||||
return tushareJSON(env), nil
|
||||
}
|
||||
|
||||
func tushareJSON(env tushareEnvelope) string {
|
||||
var buf bytes.Buffer
|
||||
enc := json.NewEncoder(&buf)
|
||||
// Error string is set via _ERROR, but if Error is empty we want the
|
||||
// key omitted. Use a typed marshal to honour the omitempty tag.
|
||||
if err := enc.Encode(env); err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"tushare: marshal result: %s"}`, err)
|
||||
}
|
||||
// json.Encoder always appends a newline — strip it.
|
||||
out := buf.Bytes()
|
||||
if n := len(out); n > 0 && out[n-1] == '\n' {
|
||||
out = out[:n-1]
|
||||
}
|
||||
return string(out)
|
||||
}
|
||||
|
||||
func tushareErrJSON(err error) string {
|
||||
return tushareJSON(tushareEnvelope{Error: err.Error()})
|
||||
}
|
||||
240
internal/agent/tool/tushare_test.go
Normal file
240
internal/agent/tool/tushare_test.go
Normal file
@@ -0,0 +1,240 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTushare_BuildRequest(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
params tushareParams
|
||||
wantToken string
|
||||
wantAPIName string
|
||||
wantFieldKey string
|
||||
wantFieldVal string
|
||||
}{
|
||||
{
|
||||
name: "minimal — token + api_name",
|
||||
params: tushareParams{
|
||||
Token: "T-abc",
|
||||
APIName: "stock_basic",
|
||||
},
|
||||
wantToken: "T-abc",
|
||||
wantAPIName: "stock_basic",
|
||||
},
|
||||
{
|
||||
name: "with params and fields",
|
||||
params: tushareParams{
|
||||
Token: "T-xyz",
|
||||
APIName: "daily",
|
||||
Params: map[string]string{"ts_code": "000001.SZ", "start_date": "20240101"},
|
||||
Fields: []string{"ts_code", "trade_date", "close"},
|
||||
},
|
||||
wantToken: "T-xyz",
|
||||
wantAPIName: "daily",
|
||||
wantFieldKey: "ts_code",
|
||||
wantFieldVal: "000001.SZ",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
body, err := buildTushareRequestBody(tc.params)
|
||||
if err != nil {
|
||||
t.Fatalf("buildTushareRequestBody: %v", err)
|
||||
}
|
||||
var got tushareRequest
|
||||
if jerr := json.Unmarshal(body, &got); jerr != nil {
|
||||
t.Fatalf("request body is not valid JSON: %v (raw=%s)", jerr, body)
|
||||
}
|
||||
if got.Token != tc.wantToken {
|
||||
t.Errorf("token = %q, want %q", got.Token, tc.wantToken)
|
||||
}
|
||||
if got.APIName != tc.wantAPIName {
|
||||
t.Errorf("api_name = %q, want %q", got.APIName, tc.wantAPIName)
|
||||
}
|
||||
if tc.wantFieldKey != "" {
|
||||
if v, ok := got.Params[tc.wantFieldKey]; !ok || v != tc.wantFieldVal {
|
||||
t.Errorf("params[%q] = %q (present=%v), want %q", tc.wantFieldKey, v, ok, tc.wantFieldVal)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTushare_ParseResponse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var (
|
||||
gotMethod string
|
||||
gotCT string
|
||||
gotBodyRaw []byte
|
||||
)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotMethod = r.Method
|
||||
gotCT = r.Header.Get("Content-Type")
|
||||
gotBodyRaw, _ = io.ReadAll(r.Body)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"code": 0,
|
||||
"msg": "",
|
||||
"data": {
|
||||
"fields": ["ts_code", "name", "industry"],
|
||||
"items": [
|
||||
["000001.SZ", "平安银行", "银行"],
|
||||
["600519.SH", "贵州茅台", "白酒"]
|
||||
]
|
||||
}
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
prev := tushareEndpoint
|
||||
tushareEndpoint = srv.URL
|
||||
t.Cleanup(func() { tushareEndpoint = prev })
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewTushareToolWith(helper)
|
||||
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"token":"T-test","api_name":"stock_basic","params":{"list_status":"L"}}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if gotMethod != http.MethodPost {
|
||||
t.Errorf("method = %q, want POST", gotMethod)
|
||||
}
|
||||
if !strings.Contains(gotCT, "application/json") {
|
||||
t.Errorf("Content-Type = %q, want application/json", gotCT)
|
||||
}
|
||||
|
||||
// Verify the request body had the right shape.
|
||||
var sentReq tushareRequest
|
||||
if jerr := json.Unmarshal(gotBodyRaw, &sentReq); jerr != nil {
|
||||
t.Fatalf("server saw malformed request body: %v (raw=%s)", jerr, gotBodyRaw)
|
||||
}
|
||||
if sentReq.Token != "T-test" {
|
||||
t.Errorf("server saw token = %q, want T-test", sentReq.Token)
|
||||
}
|
||||
if sentReq.APIName != "stock_basic" {
|
||||
t.Errorf("server saw api_name = %q, want stock_basic", sentReq.APIName)
|
||||
}
|
||||
if v := sentReq.Params["list_status"]; v != "L" {
|
||||
t.Errorf("server saw params[list_status] = %q, want L", v)
|
||||
}
|
||||
|
||||
var env tushareEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Fields) != 3 || env.Fields[0] != "ts_code" {
|
||||
t.Errorf("Fields = %v, want [ts_code name industry]", env.Fields)
|
||||
}
|
||||
if len(env.Items) != 2 {
|
||||
t.Fatalf("Items len = %d, want 2", len(env.Items))
|
||||
}
|
||||
if env.Items[0][1] != "平安银行" {
|
||||
t.Errorf("Items[0][1] = %v, want 平安银行", env.Items[0][1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestTushare_RejectsMissingToken(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewTushareTool()
|
||||
_, err := tool.InvokableRun(context.Background(),
|
||||
`{"api_name":"stock_basic"}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for missing token")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "token") {
|
||||
t.Errorf("err = %v, want to mention token", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTushare_RejectsMissingAPIName(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewTushareTool()
|
||||
_, err := tool.InvokableRun(context.Background(),
|
||||
`{"token":"T-abc"}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for missing api_name")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "api_name") {
|
||||
t.Errorf("err = %v, want to mention api_name", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTushare_UpstreamErrorCode(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"code": 40201, "msg": "权限不足"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
prev := tushareEndpoint
|
||||
tushareEndpoint = srv.URL
|
||||
t.Cleanup(func() { tushareEndpoint = prev })
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewTushareToolWith(helper)
|
||||
|
||||
_, err := tool.InvokableRun(context.Background(),
|
||||
`{"token":"T-abc","api_name":"premium_only"}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for non-zero code, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "权限不足") {
|
||||
t.Errorf("err = %v, want to surface upstream msg", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTushare_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewTushareTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "tushare" {
|
||||
t.Errorf("Name = %q, want tushare", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Tushare") {
|
||||
t.Errorf("Desc = %q, want to mention Tushare", info.Desc)
|
||||
}
|
||||
}
|
||||
128
internal/agent/tool/wencai.go
Normal file
128
internal/agent/tool/wencai.go
Normal file
@@ -0,0 +1,128 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const wencaiToolName = "wencai"
|
||||
|
||||
const wencaiToolDescription = "Query 同花顺 Wencai (问财) for natural-language stock screening " +
|
||||
"(e.g. \"近期涨停股\", \"高股息低估值\"). " +
|
||||
"STUB: Wencai has no public API; the Python implementation scrapes the " +
|
||||
"10jqka.com.cn web app. Not yet implemented in the Go Canvas. " +
|
||||
"Use the Python Canvas for Wencai queries."
|
||||
|
||||
const wencaiUnsupportedMessage = "Wencai requires web scraping of 同花顺 — not yet implemented in Go Canvas. " +
|
||||
"Use Python Canvas."
|
||||
|
||||
// wencaiParams is the JSON shape the model sends into InvokableRun.
|
||||
// The Python implementation accepts a free-form natural-language query
|
||||
// and an optional page/per-page limit. The Go stub preserves the shape
|
||||
// but rejects every invocation.
|
||||
type wencaiParams struct {
|
||||
Query string `json:"query"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PerPage int `json:"per_page,omitempty"`
|
||||
}
|
||||
|
||||
// wencaiEnvelope is the model-facing JSON shape. The stub always
|
||||
// returns a populated Error.
|
||||
type wencaiEnvelope struct {
|
||||
Items []any `json:"items,omitempty"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// WencaiTool is the Phase 3 batch 4 stub implementation of the
|
||||
// 同花顺 Wencai (问财) natural-language stock screening tool
|
||||
// (plan §2.11.4 row 19, §5 Phase 3 第 4 批).
|
||||
//
|
||||
// Wencai (https://www.iwencai.com) has no public API. The Python
|
||||
// implementation scrapes 10jqka.com.cn using session cookies and
|
||||
// reverse-engineered POST endpoints, which is fragile and legally
|
||||
// grey. A Go port would have to repeat the scraping work and the
|
||||
// reverse-engineering, and is deferred. For P3-B4 the tool is
|
||||
// registered so DSLs that reference "wencai" keep parsing, but every
|
||||
// invocation fails fast with a clear "use Python Canvas" message.
|
||||
//
|
||||
// WencaiTool does not own an HTTPHelper — it never makes network calls.
|
||||
type WencaiTool struct{}
|
||||
|
||||
// NewWencaiTool returns a WencaiTool. No HTTPHelper is allocated; the
|
||||
// stub never issues network requests.
|
||||
func NewWencaiTool() *WencaiTool { return &WencaiTool{} }
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (w *WencaiTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: wencaiToolName,
|
||||
Desc: wencaiToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Natural-language Wencai query (e.g. \"近期涨停股\", \"高股息低估值\").",
|
||||
Required: true,
|
||||
},
|
||||
"page": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Optional 1-based page number. Defaults to 1.",
|
||||
Required: false,
|
||||
},
|
||||
"per_page": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Optional results per page. Defaults to 20.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// InvokableRun validates the input shape (query is required) and
|
||||
// returns a clear "use Python Canvas" error. The model receives a
|
||||
// JSON envelope with the message in the `_ERROR` field.
|
||||
func (w *WencaiTool) InvokableRun(_ context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p wencaiParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return wencaiErrJSON(errors.New(wencaiUnsupportedMessage)),
|
||||
errors.New(wencaiUnsupportedMessage)
|
||||
}
|
||||
if p.Query == "" {
|
||||
return wencaiErrJSON(errors.New(wencaiUnsupportedMessage)),
|
||||
errors.New(wencaiUnsupportedMessage)
|
||||
}
|
||||
return wencaiErrJSON(errors.New(wencaiUnsupportedMessage)),
|
||||
errors.New(wencaiUnsupportedMessage)
|
||||
}
|
||||
|
||||
func wencaiJSON(env wencaiEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"wencai: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func wencaiErrJSON(err error) string {
|
||||
return wencaiJSON(wencaiEnvelope{Error: err.Error()})
|
||||
}
|
||||
105
internal/agent/tool/wencai_test.go
Normal file
105
internal/agent/tool/wencai_test.go
Normal file
@@ -0,0 +1,105 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestWencai_StubsUnsupported(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
args string
|
||||
}{
|
||||
{
|
||||
name: "well-formed args",
|
||||
args: `{"query":"近期涨停股","page":1,"per_page":20}`,
|
||||
},
|
||||
{
|
||||
name: "minimal args",
|
||||
args: `{"query":"高股息低估值"}`,
|
||||
},
|
||||
{
|
||||
name: "missing query",
|
||||
args: `{"page":1}`,
|
||||
},
|
||||
{
|
||||
name: "empty payload",
|
||||
args: `{}`,
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
tool := NewWencaiTool()
|
||||
out, err := tool.InvokableRun(context.Background(), tc.args)
|
||||
if err == nil {
|
||||
t.Fatalf("expected error, got nil (out=%s)", out)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "同花顺") {
|
||||
t.Errorf("err = %q, want to mention 同花顺", err.Error())
|
||||
}
|
||||
var env wencaiEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error == "" {
|
||||
t.Errorf("env.Error = empty, want populated")
|
||||
}
|
||||
if !strings.Contains(env.Error, "同花顺") {
|
||||
t.Errorf("env.Error = %q, want to mention 同花顺", env.Error)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_RejectsMalformedJSON(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewWencaiTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{not json`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for malformed JSON, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "同花顺") {
|
||||
t.Errorf("err = %q, want to mention 同花顺", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewWencaiTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "wencai" {
|
||||
t.Errorf("Name = %q, want wencai", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Wencai") {
|
||||
t.Errorf("Desc = %q, want to mention Wencai", info.Desc)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "STUB") && !strings.Contains(info.Desc, "Python") {
|
||||
t.Errorf("Desc = %q, want to flag stub status", info.Desc)
|
||||
}
|
||||
}
|
||||
186
internal/agent/tool/wikipedia.go
Normal file
186
internal/agent/tool/wikipedia.go
Normal file
@@ -0,0 +1,186 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const wikipediaToolName = "wikipedia"
|
||||
|
||||
const wikipediaToolDescription = "Search Wikipedia and return matching articles as {title, snippet, url}."
|
||||
|
||||
// wikipediaParams is the JSON shape the model sends into InvokableRun.
|
||||
// lang is the language subdomain (e.g. "en", "zh", "de"); max_results
|
||||
// defaults to 5 when unset or non-positive.
|
||||
type wikipediaParams struct {
|
||||
Query string `json:"query"`
|
||||
Lang string `json:"lang"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
// wikipediaResult is one row in the upstream `query.search` array.
|
||||
type wikipediaResult struct {
|
||||
Title string `json:"title"`
|
||||
Snippet string `json:"snippet"`
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
// wikipediaResponse is the upstream MediaWiki API envelope.
|
||||
type wikipediaResponse struct {
|
||||
Query struct {
|
||||
Search []wikipediaResult `json:"search"`
|
||||
} `json:"query"`
|
||||
}
|
||||
|
||||
// wikipediaEnvelope is what the model sees. It mirrors the Python tool's
|
||||
// output: a flat list of {title, snippet, url} entries.
|
||||
type wikipediaEnvelope struct {
|
||||
Results []wikipediaResult `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// WikipediaTool is the Phase 3 batch 2 implementation of the Wikipedia
|
||||
// search tool (plan §2.11.4 row 20, §5 Phase 3 第 2 批). It calls the
|
||||
// public MediaWiki action API via the shared HTTPHelper and returns the
|
||||
// top N matches for the query.
|
||||
type WikipediaTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewWikipediaTool returns a WikipediaTool using the default HTTPHelper.
|
||||
func NewWikipediaTool() *WikipediaTool {
|
||||
return NewWikipediaToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewWikipediaToolWith returns a WikipediaTool that uses the provided
|
||||
// HTTPHelper. Useful for tests that want to inject a custom transport.
|
||||
func NewWikipediaToolWith(h *HTTPHelper) *WikipediaTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &WikipediaTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (w *WikipediaTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: wikipediaToolName,
|
||||
Desc: wikipediaToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query",
|
||||
Required: true,
|
||||
},
|
||||
"lang": {
|
||||
Type: schema.String,
|
||||
Desc: `Language subdomain. Defaults to "en".`,
|
||||
Required: false,
|
||||
},
|
||||
"max_results": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of results to return. Defaults to 5.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildWikipediaURL constructs the MediaWiki action=query URL. Centralized
|
||||
// so the test suite can verify URL encoding without spinning up a server.
|
||||
func buildWikipediaURL(lang, query string, maxResults int) string {
|
||||
if lang == "" {
|
||||
lang = "en"
|
||||
}
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
return fmt.Sprintf(
|
||||
"https://%s.wikipedia.org/w/api.php?action=query&list=search&format=json&srlimit=%d&srsearch=%s",
|
||||
lang,
|
||||
maxResults,
|
||||
url.QueryEscape(query),
|
||||
)
|
||||
}
|
||||
|
||||
// InvokableRun performs the Wikipedia search.
|
||||
func (w *WikipediaTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p wikipediaParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return wikipediaErrJSON(fmt.Errorf("wikipedia: parse arguments: %w", err)),
|
||||
fmt.Errorf("wikipedia: parse arguments: %w", err)
|
||||
}
|
||||
if p.Query == "" {
|
||||
return wikipediaErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("wikipedia: query is required")
|
||||
}
|
||||
|
||||
endpoint := buildWikipediaURL(p.Lang, p.Query, p.MaxResults)
|
||||
resp, err := w.helper.Do(ctx, http.MethodGet, endpoint, "", "", nil)
|
||||
if err != nil {
|
||||
return wikipediaErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return wikipediaErrJSON(fmt.Errorf("wikipedia: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("wikipedia: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw wikipediaResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return wikipediaErrJSON(fmt.Errorf("wikipedia: decode response: %w", err)),
|
||||
fmt.Errorf("wikipedia: decode response: %w", err)
|
||||
}
|
||||
|
||||
// Compose canonical Wikipedia URLs for the snippets. The MediaWiki
|
||||
// search endpoint doesn't include a URL; the conventional one is
|
||||
// https://<lang>.wikipedia.org/wiki/<Title>.
|
||||
lang := p.Lang
|
||||
if lang == "" {
|
||||
lang = "en"
|
||||
}
|
||||
results := make([]wikipediaResult, 0, len(raw.Query.Search))
|
||||
for _, r := range raw.Query.Search {
|
||||
results = append(results, wikipediaResult{
|
||||
Title: r.Title,
|
||||
Snippet: r.Snippet,
|
||||
URL: fmt.Sprintf("https://%s.wikipedia.org/wiki/%s", lang, url.PathEscape(r.Title)),
|
||||
})
|
||||
}
|
||||
return wikipediaJSON(wikipediaEnvelope{Results: results}), nil
|
||||
}
|
||||
|
||||
func wikipediaJSON(env wikipediaEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"wikipedia: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func wikipediaErrJSON(err error) string {
|
||||
return wikipediaJSON(wikipediaEnvelope{Error: err.Error()})
|
||||
}
|
||||
194
internal/agent/tool/wikipedia_test.go
Normal file
194
internal/agent/tool/wikipedia_test.go
Normal file
@@ -0,0 +1,194 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestWikipedia_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
lang string
|
||||
query string
|
||||
max int
|
||||
wantHost string
|
||||
wantPath string
|
||||
}{
|
||||
{
|
||||
name: "en default",
|
||||
lang: "",
|
||||
query: "rag flow",
|
||||
max: 0,
|
||||
wantHost: "en.wikipedia.org",
|
||||
wantPath: "/w/api.php",
|
||||
},
|
||||
{
|
||||
name: "de explicit",
|
||||
lang: "de",
|
||||
query: "Berlin",
|
||||
max: 3,
|
||||
wantHost: "de.wikipedia.org",
|
||||
wantPath: "/w/api.php",
|
||||
},
|
||||
{
|
||||
name: "spaces encoded",
|
||||
lang: "en",
|
||||
query: "a b c",
|
||||
max: 1,
|
||||
wantHost: "en.wikipedia.org",
|
||||
wantPath: "/w/api.php",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildWikipediaURL(tc.lang, tc.query, tc.max)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != tc.wantHost {
|
||||
t.Errorf("host = %q, want %q", u.Host, tc.wantHost)
|
||||
}
|
||||
if u.Path != tc.wantPath {
|
||||
t.Errorf("path = %q, want %q", u.Path, tc.wantPath)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("action") != "query" {
|
||||
t.Errorf("action = %q, want query", q.Get("action"))
|
||||
}
|
||||
if q.Get("list") != "search" {
|
||||
t.Errorf("list = %q, want search", q.Get("list"))
|
||||
}
|
||||
if q.Get("format") != "json" {
|
||||
t.Errorf("format = %q, want json", q.Get("format"))
|
||||
}
|
||||
if q.Get("srsearch") != tc.query {
|
||||
t.Errorf("srsearch = %q, want %q (raw query, not pre-encoded)", q.Get("srsearch"), tc.query)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWikipedia_ParseResults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"query": {
|
||||
"search": [
|
||||
{"title":"RAG","snippet":"<span>rag</span> is ..."},
|
||||
{"title":"Retrieval-augmented generation","snippet":"<b>RAG</b> is ..."}
|
||||
]
|
||||
}
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
// Point the hard-coded en.wikipedia.org endpoint at the test server
|
||||
// by injecting a transport that rewrites the request host.
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewWikipediaToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(), `{"query":"RAG","lang":"en","max_results":5}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var env wikipediaEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
}
|
||||
if env.Results[0].Title != "RAG" {
|
||||
t.Errorf("Results[0].Title = %q, want RAG", env.Results[0].Title)
|
||||
}
|
||||
if !strings.HasPrefix(env.Results[0].URL, "https://en.wikipedia.org/wiki/") {
|
||||
t.Errorf("Results[0].URL = %q, want to start with https://en.wikipedia.org/wiki/", env.Results[0].URL)
|
||||
}
|
||||
}
|
||||
|
||||
// rewriteHostTransport returns a RoundTripper that rewrites the request
|
||||
// host to the given test server URL. Used to point the hard-coded
|
||||
// en.wikipedia.org endpoint at a httptest.Server without changing the
|
||||
// production URL builder.
|
||||
func rewriteHostTransport(srvURL string) http.RoundTripper {
|
||||
u, err := url.Parse(srvURL)
|
||||
if err != nil {
|
||||
panic("rewriteHostTransport: bad srvURL: " + err.Error())
|
||||
}
|
||||
return &hostSwapRT{inner: http.DefaultTransport, host: u.Host, scheme: u.Scheme}
|
||||
}
|
||||
|
||||
type hostSwapRT struct {
|
||||
inner http.RoundTripper
|
||||
host string
|
||||
scheme string
|
||||
}
|
||||
|
||||
func (t *hostSwapRT) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
r2 := req.Clone(req.Context())
|
||||
r2.URL.Scheme = t.scheme
|
||||
r2.URL.Host = t.host
|
||||
r2.Host = t.host
|
||||
return t.inner.RoundTrip(r2)
|
||||
}
|
||||
|
||||
func TestWikipedia_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewWikipediaTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "wikipedia" {
|
||||
t.Errorf("Name = %q, want wikipedia", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Wikipedia") {
|
||||
t.Errorf("Desc = %q, want to mention Wikipedia", info.Desc)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWikipedia_RequiresQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewWikipediaTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{"query":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty query")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "query") {
|
||||
t.Errorf("err = %v, want to mention query", err)
|
||||
}
|
||||
}
|
||||
172
internal/agent/tool/yahoo_finance.go
Normal file
172
internal/agent/tool/yahoo_finance.go
Normal file
@@ -0,0 +1,172 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const yahooFinanceToolName = "yahoo_finance"
|
||||
|
||||
const yahooFinanceToolDescription = "Fetch stock quote snapshots from Yahoo Finance. Returns quoteResponse.result[].{symbol, regularMarketPrice, currency, regularMarketChangePercent}."
|
||||
|
||||
// yahooFinanceParams is the JSON shape the model sends into InvokableRun.
|
||||
type yahooFinanceParams struct {
|
||||
Symbols []string `json:"symbols"`
|
||||
Fields []string `json:"fields"`
|
||||
}
|
||||
|
||||
// yahooFinanceQuote is one element of the upstream result array.
|
||||
type yahooFinanceQuote struct {
|
||||
Symbol string `json:"symbol"`
|
||||
RegularMarketPrice float64 `json:"regularMarketPrice"`
|
||||
Currency string `json:"currency"`
|
||||
RegularMarketChangePercent float64 `json:"regularMarketChangePercent"`
|
||||
}
|
||||
|
||||
// yahooFinanceResponse is the upstream Yahoo Finance /v7/finance/quote
|
||||
// envelope.
|
||||
type yahooFinanceResponse struct {
|
||||
QuoteResponse struct {
|
||||
Result []yahooFinanceQuote `json:"result"`
|
||||
Error any `json:"error,omitempty"`
|
||||
} `json:"quoteResponse"`
|
||||
}
|
||||
|
||||
// yahooFinanceEnvelope is what the model sees.
|
||||
type yahooFinanceEnvelope struct {
|
||||
Results []yahooFinanceQuote `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// yahooFinanceEndpoint is the Yahoo Finance quote URL. Exposed as a
|
||||
// package var so tests can substitute a httptest.Server URL.
|
||||
var yahooFinanceEndpoint = "https://query1.finance.yahoo.com/v7/finance/quote"
|
||||
|
||||
// YahooFinanceTool is the Phase 3 batch 3 implementation of the
|
||||
// Yahoo Finance quote tool (plan §2.11.4 row 22, §5 Phase 3 第 3 批).
|
||||
// It performs an unauthenticated GET against the public quote API
|
||||
// via the shared HTTPHelper and returns the parsed quote records.
|
||||
type YahooFinanceTool struct {
|
||||
helper *HTTPHelper
|
||||
}
|
||||
|
||||
// NewYahooFinanceTool returns a YahooFinanceTool using the default
|
||||
// HTTPHelper.
|
||||
func NewYahooFinanceTool() *YahooFinanceTool {
|
||||
return NewYahooFinanceToolWith(NewHTTPHelper())
|
||||
}
|
||||
|
||||
// NewYahooFinanceToolWith returns a YahooFinanceTool that uses the
|
||||
// provided HTTPHelper. Useful for tests.
|
||||
func NewYahooFinanceToolWith(h *HTTPHelper) *YahooFinanceTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &YahooFinanceTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func (y *YahooFinanceTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: yahooFinanceToolName,
|
||||
Desc: yahooFinanceToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"symbols": {
|
||||
Type: schema.Array,
|
||||
Desc: "Stock symbols to look up (e.g. AAPL, MSFT, 0005.HK).",
|
||||
Required: true,
|
||||
},
|
||||
"fields": {
|
||||
Type: schema.Array,
|
||||
Desc: "Optional list of fields to request via the `fields` query parameter.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildYahooFinanceURL composes the quote URL with the symbol list
|
||||
// and an optional `fields` parameter. Centralized for testability.
|
||||
func buildYahooFinanceURL(symbols []string, fields []string) string {
|
||||
q := url.Values{}
|
||||
q.Set("symbols", strings.Join(symbols, ","))
|
||||
if len(fields) > 0 {
|
||||
q.Set("fields", strings.Join(fields, ","))
|
||||
}
|
||||
return yahooFinanceEndpoint + "?" + q.Encode()
|
||||
}
|
||||
|
||||
// InvokableRun performs the Yahoo Finance quote lookup.
|
||||
func (y *YahooFinanceTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var p yahooFinanceParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
||||
return yahooFinanceErrJSON(fmt.Errorf("yahoo_finance: parse arguments: %w", err)),
|
||||
fmt.Errorf("yahoo_finance: parse arguments: %w", err)
|
||||
}
|
||||
if len(p.Symbols) == 0 {
|
||||
return yahooFinanceErrJSON(fmt.Errorf("symbols is required and must be non-empty")),
|
||||
fmt.Errorf("yahoo_finance: symbols is required and must be non-empty")
|
||||
}
|
||||
|
||||
endpoint := buildYahooFinanceURL(p.Symbols, p.Fields)
|
||||
// Yahoo Finance returns 401 unless we send a User-Agent that
|
||||
// looks like a real browser. curl-style UA is the conventional
|
||||
// workaround for the public (unauthenticated) endpoint.
|
||||
headers := map[string]string{
|
||||
"User-Agent": "Mozilla/5.0 (compatible; ragflow/1.0)",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
|
||||
resp, err := y.helper.Do(ctx, http.MethodGet, endpoint, "", "", headers)
|
||||
if err != nil {
|
||||
return yahooFinanceErrJSON(err), err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return yahooFinanceErrJSON(fmt.Errorf("yahoo_finance: upstream returned %d", resp.StatusCode)),
|
||||
fmt.Errorf("yahoo_finance: upstream returned %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var raw yahooFinanceResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
|
||||
return yahooFinanceErrJSON(fmt.Errorf("yahoo_finance: decode response: %w", err)),
|
||||
fmt.Errorf("yahoo_finance: decode response: %w", err)
|
||||
}
|
||||
return yahooFinanceJSON(yahooFinanceEnvelope{Results: raw.QuoteResponse.Result}), nil
|
||||
}
|
||||
|
||||
func yahooFinanceJSON(env yahooFinanceEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"yahoo_finance: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func yahooFinanceErrJSON(err error) string {
|
||||
return yahooFinanceJSON(yahooFinanceEnvelope{Error: err.Error()})
|
||||
}
|
||||
160
internal/agent/tool/yahoo_finance_test.go
Normal file
160
internal/agent/tool/yahoo_finance_test.go
Normal file
@@ -0,0 +1,160 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestYahooFinance_BuildURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
symbols []string
|
||||
fields []string
|
||||
wantSymbols string
|
||||
wantFields string
|
||||
wantHost string
|
||||
}{
|
||||
{
|
||||
name: "single symbol, no fields",
|
||||
symbols: []string{"AAPL"},
|
||||
fields: nil,
|
||||
wantSymbols: "AAPL",
|
||||
wantHost: "query1.finance.yahoo.com",
|
||||
},
|
||||
{
|
||||
name: "multi symbol, with fields",
|
||||
symbols: []string{"AAPL", "MSFT", "0005.HK"},
|
||||
fields: []string{"symbol", "regularMarketPrice"},
|
||||
wantSymbols: "AAPL,MSFT,0005.HK",
|
||||
wantFields: "symbol,regularMarketPrice",
|
||||
wantHost: "query1.finance.yahoo.com",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := buildYahooFinanceURL(tc.symbols, tc.fields)
|
||||
u, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
if u.Host != tc.wantHost {
|
||||
t.Errorf("host = %q, want %q", u.Host, tc.wantHost)
|
||||
}
|
||||
if u.Path != "/v7/finance/quote" {
|
||||
t.Errorf("path = %q, want /v7/finance/quote", u.Path)
|
||||
}
|
||||
q := u.Query()
|
||||
if q.Get("symbols") != tc.wantSymbols {
|
||||
t.Errorf("symbols = %q, want %q", q.Get("symbols"), tc.wantSymbols)
|
||||
}
|
||||
if tc.wantFields != "" && q.Get("fields") != tc.wantFields {
|
||||
t.Errorf("fields = %q, want %q", q.Get("fields"), tc.wantFields)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestYahooFinance_ParseQuote(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var gotUA string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotUA = r.Header.Get("User-Agent")
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"quoteResponse": {
|
||||
"result": [
|
||||
{"symbol":"AAPL","regularMarketPrice":189.5,"currency":"USD","regularMarketChangePercent":1.23},
|
||||
{"symbol":"MSFT","regularMarketPrice":421.0,"currency":"USD","regularMarketChangePercent":-0.5}
|
||||
]
|
||||
}
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewYahooFinanceToolWith(helper)
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"symbols":["AAPL","MSFT"]}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if !strings.Contains(gotUA, "ragflow") {
|
||||
t.Errorf("User-Agent = %q, want to contain ragflow", gotUA)
|
||||
}
|
||||
|
||||
var env yahooFinanceEnvelope
|
||||
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
}
|
||||
if env.Results[0].Symbol != "AAPL" {
|
||||
t.Errorf("Results[0].Symbol = %q, want AAPL", env.Results[0].Symbol)
|
||||
}
|
||||
if env.Results[0].RegularMarketPrice != 189.5 {
|
||||
t.Errorf("Results[0].Price = %v, want 189.5", env.Results[0].RegularMarketPrice)
|
||||
}
|
||||
if env.Results[1].RegularMarketChangePercent != -0.5 {
|
||||
t.Errorf("Results[1].ChangePct = %v, want -0.5", env.Results[1].RegularMarketChangePercent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestYahooFinance_RequiresSymbols(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewYahooFinanceTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{"symbols":[]}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty symbols")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "symbols") {
|
||||
t.Errorf("err = %v, want to mention symbols", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestYahooFinance_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewYahooFinanceTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "yahoo_finance" {
|
||||
t.Errorf("Name = %q, want yahoo_finance", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "Yahoo") {
|
||||
t.Errorf("Desc = %q, want to mention Yahoo", info.Desc)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user