mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-11-11 13:51:01 +08:00
fix(agent/tool): read the Tushare token from node config over https and drop it from the model schema (#21679)
This commit is contained in:
4 files changed
+159
-18
No files matched your search
@@ -69,7 +69,7 @@ var registry = map[string]Factory{
|
||||
"sofya": buildSofyaTool,
|
||||
"tavily": buildTavilyTool,
|
||||
"tavily_extract": buildTavilyExtractTool,
|
||||
"tushare": noConfig("tushare", func() einotool.BaseTool { return NewTushareTool() }),
|
||||
"tushare": buildTushareTool,
|
||||
"wencai": buildWencaiTool,
|
||||
"web_crawler": noConfig("web_crawler", func() einotool.BaseTool { return NewCrawlerTool() }),
|
||||
"wikipedia": buildWikipediaTool,
|
||||
@@ -869,6 +869,20 @@ func buildSearch1APICrawlTool(params map[string]any) (einotool.BaseTool, error)
|
||||
return newSearch1APICrawlTool(nil, apiKey, defaultURL), nil
|
||||
}
|
||||
|
||||
func buildTushareTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
var token string
|
||||
if value, exists := params["token"]; exists {
|
||||
var valid bool
|
||||
token, valid = value.(string)
|
||||
if !valid {
|
||||
return nil, fmt.Errorf("agent tool: tool %q requires string node-level param token", tushareToolName)
|
||||
}
|
||||
}
|
||||
// An empty token builds fine; InvokableRun reports it as a node
|
||||
// configuration error, so the registry can still construct every tool.
|
||||
return newTushareTool(nil, tushareParams{Token: strings.TrimSpace(token)}), nil
|
||||
}
|
||||
|
||||
func buildWikipediaTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
topN := defaultWikipediaTopN
|
||||
language := defaultWikipediaLanguage
|
||||
|
||||
@@ -233,6 +233,35 @@ func TestBuildAll_AllRegisteredTools(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildTushareTool_NodeParams pins the node-level credential channel:
|
||||
// the token comes from tool_params, an empty set still builds (Invoke reports
|
||||
// the missing configuration), and a non-string token is rejected.
|
||||
func TestBuildTushareTool_NodeParams(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
built, err := BuildByName("tushare", map[string]any{"token": "T-node"})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildByName(tushare) with token: %v", err)
|
||||
}
|
||||
tool, ok := built.(*TushareTool)
|
||||
if !ok {
|
||||
t.Fatalf("tool type = %T, want *TushareTool", built)
|
||||
}
|
||||
if tool.defaults.Token != "T-node" {
|
||||
t.Fatalf("defaults.Token = %q, want T-node", tool.defaults.Token)
|
||||
}
|
||||
|
||||
if _, err = BuildByName("tushare", map[string]any{}); err != nil {
|
||||
t.Fatalf("BuildByName(tushare) without params: %v", err)
|
||||
}
|
||||
|
||||
if _, err = BuildByName("tushare", map[string]any{"token": 1}); err == nil {
|
||||
t.Fatal("expected a type error for a non-string token")
|
||||
} else if !strings.Contains(err.Error(), "string node-level param token") {
|
||||
t.Fatalf("err = %q, want the node-level token type error", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAll_ExeSQLRequiresNodeParams(t *testing.T) {
|
||||
_, err := BuildAll([]string{"execute_sql"}, nil)
|
||||
if err == nil {
|
||||
|
||||
@@ -34,13 +34,17 @@ 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"
|
||||
// tushareEndpoint is the Tushare Pro REST endpoint. The API only reads the
|
||||
// request body, so plaintext HTTP would ship the account token in the clear;
|
||||
// keep it on TLS. Exposed as a package var so tests can substitute a
|
||||
// httptest.Server URL.
|
||||
var tushareEndpoint = "https://api.tushare.pro"
|
||||
|
||||
// tushareParams is the JSON shape the model sends into InvokableRun.
|
||||
//
|
||||
// - Token (required): the Tushare Pro API token (积分账户).
|
||||
// - Token: the Tushare Pro API token (积分账户). Resolved from the Canvas
|
||||
// node's tool_params, never from the model; a model-supplied value is
|
||||
// ignored so the credential cannot travel through tool-call arguments.
|
||||
// - 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,
|
||||
@@ -101,21 +105,27 @@ type tushareEnvelope struct {
|
||||
// TushareTool uses the shared HTTPHelper for retry/timeout/OTel
|
||||
// propagation.
|
||||
type TushareTool struct {
|
||||
helper *HTTPHelper
|
||||
helper *HTTPHelper
|
||||
defaults tushareParams
|
||||
}
|
||||
|
||||
// NewTushareTool returns a TushareTool using the default HTTPHelper.
|
||||
func NewTushareTool() *TushareTool {
|
||||
return NewTushareToolWith(NewHTTPHelper())
|
||||
return newTushareTool(NewHTTPHelper(), tushareParams{})
|
||||
}
|
||||
|
||||
// NewTushareToolWith returns a TushareTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewTushareToolWith(h *HTTPHelper) *TushareTool {
|
||||
return newTushareTool(h, tushareParams{})
|
||||
}
|
||||
|
||||
// newTushareTool carries the node-level defaults, including the API token.
|
||||
func newTushareTool(h *HTTPHelper, defaults tushareParams) *TushareTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
}
|
||||
return &TushareTool{helper: h}
|
||||
return &TushareTool{helper: h, defaults: defaults}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
@@ -124,11 +134,6 @@ func (t *TushareTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
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.",
|
||||
@@ -178,9 +183,12 @@ func (t *TushareTool) InvokableRun(ctx context.Context, argsJSON string, _ ...to
|
||||
return tushareErrJSON(fmt.Errorf("tushare: parse arguments: %w", err)),
|
||||
fmt.Errorf("tushare: parse arguments: %w", err)
|
||||
}
|
||||
// The token is a node-level credential: the node's value overrides whatever
|
||||
// the model sent, so it never has to travel through tool-call arguments.
|
||||
p.Token = t.defaults.Token
|
||||
if p.Token == "" {
|
||||
return tushareErrJSON(fmt.Errorf("tushare: token is required")),
|
||||
fmt.Errorf("tushare: token is required")
|
||||
return tushareErrJSON(fmt.Errorf("tushare: token is not configured on the node")),
|
||||
fmt.Errorf("tushare: token is not configured on the node")
|
||||
}
|
||||
if p.APIName == "" {
|
||||
return tushareErrJSON(fmt.Errorf("tushare: api_name is required")),
|
||||
|
||||
@@ -21,10 +21,15 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// defaultTushareEndpoint snapshots the package endpoint before any test swaps
|
||||
// it out, so the HTTPS assertion is not affected by parallel test mutations.
|
||||
var defaultTushareEndpoint = tushareEndpoint
|
||||
|
||||
func TestTushare_BuildRequest(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -120,7 +125,7 @@ func TestTushare_ParseResponse(t *testing.T) {
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewTushareToolWith(helper)
|
||||
tool := newTushareTool(helper, tushareParams{Token: "T-test"})
|
||||
|
||||
out, err := tool.InvokableRun(ctx,
|
||||
`{"token":"T-test","api_name":"stock_basic","params":{"list_status":"L"}}`)
|
||||
@@ -180,13 +185,68 @@ func TestTushare_RejectsMissingToken(t *testing.T) {
|
||||
if !strings.Contains(err.Error(), "token") {
|
||||
t.Errorf("err = %v, want to mention token", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not configured on the node") {
|
||||
t.Errorf("err = %v, want to name the node-level configuration", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTushare_ModelSuppliedTokenIsIgnored pins the credential channel: the
|
||||
// token comes from the node's tool_params and a model-supplied value must not
|
||||
// reach upstream.
|
||||
func TestTushare_ModelSuppliedTokenIsIgnored(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := t.Context()
|
||||
|
||||
var gotBodyRaw []byte
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotBodyRaw, _ = io.ReadAll(r.Body)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"code":0,"data":{"fields":[],"items":[]}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
prev := tushareEndpoint
|
||||
tushareEndpoint = srv.URL
|
||||
t.Cleanup(func() { tushareEndpoint = prev })
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := newTushareTool(helper, tushareParams{Token: "node-token"})
|
||||
|
||||
if _, err := tool.InvokableRun(ctx,
|
||||
`{"token":"model-token","api_name":"stock_basic"}`); err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
|
||||
var sentReq tushareRequest
|
||||
if err := json.Unmarshal(gotBodyRaw, &sentReq); err != nil {
|
||||
t.Fatalf("server saw malformed request body: %v (raw=%s)", err, gotBodyRaw)
|
||||
}
|
||||
if sentReq.Token != "node-token" {
|
||||
t.Errorf("server saw token = %q, want the node-level node-token", sentReq.Token)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTushare_DefaultEndpointIsHTTPS(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Other tests substitute tushareEndpoint with an httptest URL, so compare
|
||||
// against the package default captured before any test could mutate it.
|
||||
parsed, err := url.Parse(defaultTushareEndpoint)
|
||||
if err != nil {
|
||||
t.Fatalf("parse tushareEndpoint: %v", err)
|
||||
}
|
||||
if parsed.Scheme != "https" {
|
||||
t.Errorf("tushareEndpoint scheme = %q, want https", parsed.Scheme)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTushare_RejectsMissingAPIName(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := t.Context()
|
||||
|
||||
tool := NewTushareTool()
|
||||
tool := newTushareTool(nil, tushareParams{Token: "T-abc"})
|
||||
_, err := tool.InvokableRun(ctx,
|
||||
`{"token":"T-abc"}`)
|
||||
if err == nil {
|
||||
@@ -214,7 +274,7 @@ func TestTushare_UpstreamErrorCode(t *testing.T) {
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{
|
||||
Transport: rewriteHostTransport(srv.URL),
|
||||
})
|
||||
tool := NewTushareToolWith(helper)
|
||||
tool := newTushareTool(helper, tushareParams{Token: "T-abc"})
|
||||
|
||||
_, err := tool.InvokableRun(ctx,
|
||||
`{"token":"T-abc","api_name":"premium_only"}`)
|
||||
@@ -241,4 +301,34 @@ func TestTushare_Info(t *testing.T) {
|
||||
if !strings.Contains(info.Desc, "Tushare") {
|
||||
t.Errorf("Desc = %q, want to mention Tushare", info.Desc)
|
||||
}
|
||||
schema, err := info.ParamsOneOf.ToJSONSchema()
|
||||
if err != nil {
|
||||
t.Fatalf("Info schema: %v", err)
|
||||
}
|
||||
encoded, err := json.Marshal(schema)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal Info schema: %v", err)
|
||||
}
|
||||
var decoded struct {
|
||||
Properties map[string]any `json:"properties"`
|
||||
Required []string `json:"required"`
|
||||
}
|
||||
if err = json.Unmarshal(encoded, &decoded); err != nil {
|
||||
t.Fatalf("decode Info schema: %v", err)
|
||||
}
|
||||
if _, hasToken := decoded.Properties["token"]; hasToken {
|
||||
t.Error("schema still exposes token; the credential must come from the node config")
|
||||
}
|
||||
if _, ok := decoded.Properties["api_name"]; !ok {
|
||||
t.Error("schema must keep api_name")
|
||||
}
|
||||
found := false
|
||||
for _, name := range decoded.Required {
|
||||
if name == "api_name" {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("schema required = %v, want api_name", decoded.Required)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user