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:
Zhichang Yu
2026-06-12 22:58:28 +08:00
committed by GitHub
parent cafa0f2e4f
commit 3fa15c0e2f
232 changed files with 44641 additions and 3993 deletions

View 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()})
}

View 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)
}
}

View 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()})
}

View 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)
}
}

View 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 ""
}

View 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)
}
}

View 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
}

View 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).

View 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()})
}

View 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)
}
}

View 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()})
}

View 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")
}
}

View 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()})
}

View 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)
}
}

View 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()
}

View 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)
}
}

View 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()})
}

View 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)
}
}

View 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()})
}

View 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()})
}

View 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)
}
}

View 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)
}
}

View 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
}

View 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,
}
}

View 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()})
}

View 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)
}
}

View 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), &params); 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()})
}

View 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)
}
}

View 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)
}

View 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)
}
}

View 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
}
}

View 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
}

View 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)
}

View 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)
}
}

View 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()})
}

View 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
View 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()
}

View 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)
}
})
}
}

View 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()})
}

View 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)
}
}

View 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()})
}

View 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)
}
}

View 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()})
}

View 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)
}
}

View 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()})
}

View 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)
}
}

View 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()})
}

View 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)
}
}