mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-16 12:47:19 +08:00
fix(go-agent): Align Go WenCai and SearXNG agent components with python (#16854)
## Summary - Align Go WenCai and SearXNG behavior, schemas, and node parameters with Python. - Add the `WenCai` and `SearXNG` Canvas components and register their tool factories. - Match Python's current WenCai behavior by returning an empty report while its upstream request is disabled. - Add SearXNG request validation, SSRF-safe DNS pinning, raw result preservation, and reference rendering. - Support context cancellation, error envelopes, and lock-safe retrieval references. ## Tests Passed: - `bash build.sh --test ./internal/agent/tool/...` - `bash build.sh --test ./internal/agent/component/...` - `bash build.sh --test ./internal/agent/runtime/...` - `bash build.sh --test ./internal/agent/...` - `cd web && npm run type-check` <img width="1900" height="1102" alt="image" src="https://github.com/user-attachments/assets/ec77d217-d9fd-455a-96ec-9aabf6841109" /> <img width="1900" height="1102" alt="image" src="https://github.com/user-attachments/assets/52ac129f-cb65-453d-ae48-cc518803ac23" />
This commit is contained in:
252
internal/agent/component/searxng.go
Normal file
252
internal/agent/component/searxng.go
Normal file
@@ -0,0 +1,252 @@
|
||||
//
|
||||
// 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 component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha1"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
agenttool "ragflow/internal/agent/tool"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
const searxngPromptTokenLimit = 200000
|
||||
|
||||
var searxngDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`)
|
||||
|
||||
var searxngNewlinePattern = regexp.MustCompile(`\n+`)
|
||||
|
||||
type searxngInvoker interface {
|
||||
InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error)
|
||||
}
|
||||
|
||||
type searxngComponent struct {
|
||||
inner searxngInvoker
|
||||
}
|
||||
|
||||
func newSearXNGComponent(params map[string]any) (Component, error) {
|
||||
toolParams := make(map[string]any, 2)
|
||||
for _, key := range []string{"top_n", "searxng_url"} {
|
||||
if value, ok := params[key]; ok {
|
||||
toolParams[key] = value
|
||||
}
|
||||
}
|
||||
inner, err := agenttool.BuildByName("searxng", toolParams)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
invoker, ok := inner.(searxngInvoker)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("SearXNG: tool does not implement InvokableRun")
|
||||
}
|
||||
return newSearXNGComponentWithInvoker(invoker), nil
|
||||
}
|
||||
|
||||
func newSearXNGComponentWithInvoker(inner searxngInvoker) Component {
|
||||
return &searxngComponent{inner: inner}
|
||||
}
|
||||
|
||||
func (c *searxngComponent) Name() string { return "SearXNG" }
|
||||
|
||||
func (c *searxngComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"query": "The search keywords to execute with SearXNG.",
|
||||
"searxng_url": "The base URL of the SearXNG instance.",
|
||||
}
|
||||
}
|
||||
|
||||
func (c *searxngComponent) Outputs() map[string]string {
|
||||
return map[string]string{
|
||||
"formalized_content": "Rendered SearXNG references for downstream LLM prompts.",
|
||||
"json": "Raw SearXNG result list.",
|
||||
}
|
||||
}
|
||||
|
||||
func (c *searxngComponent) GetInputForm() map[string]any {
|
||||
return map[string]any{
|
||||
"query": map[string]any{
|
||||
"name": "Query",
|
||||
"type": "line",
|
||||
},
|
||||
"searxng_url": map[string]any{
|
||||
"name": "SearXNG URL",
|
||||
"type": "line",
|
||||
"placeholder": "http://localhost:4000",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *searxngComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
query := stringParam(inputs["query"])
|
||||
if strings.TrimSpace(query) == "" {
|
||||
return map[string]any{"formalized_content": "", "json": []any{}}, nil
|
||||
}
|
||||
args := map[string]any{"query": query}
|
||||
if searxngURL, ok := inputs["searxng_url"].(string); ok {
|
||||
args["searxng_url"] = searxngURL
|
||||
}
|
||||
argsJSON, err := json.Marshal(args)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("canvas: SearXNG: encode inputs: %w", err)
|
||||
}
|
||||
|
||||
out, invokeErr := c.inner.InvokableRun(ctx, string(argsJSON))
|
||||
decoded := parseToolEnvelope(out)
|
||||
results := anySlice(decoded["results"])
|
||||
if message, _ := decoded["_ERROR"].(string); strings.TrimSpace(message) != "" {
|
||||
return map[string]any{
|
||||
"formalized_content": "",
|
||||
"json": results,
|
||||
"_ERROR": message,
|
||||
}, nil
|
||||
}
|
||||
if invokeErr != nil {
|
||||
return nil, fmt.Errorf("canvas: SearXNG: %w", invokeErr)
|
||||
}
|
||||
|
||||
chunks, docAggs := buildSearXNGReferences(results)
|
||||
if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil {
|
||||
state.SetRetrievalReferences(chunks, docAggs)
|
||||
}
|
||||
return map[string]any{
|
||||
"formalized_content": renderSearXNGReferences(chunks, searxngPromptTokenLimit),
|
||||
"json": results,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *searxngComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func buildSearXNGReferences(results []any) ([]map[string]any, []map[string]any) {
|
||||
chunks := make([]map[string]any, 0, len(results))
|
||||
docAggs := make([]map[string]any, 0, len(results))
|
||||
for _, result := range results {
|
||||
item, ok := result.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
content, _ := item["content"].(string)
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
content = searxngDataImagePattern.ReplaceAllString(content, "")
|
||||
runes := []rune(content)
|
||||
if len(runes) > 10000 {
|
||||
content = string(runes[:10000])
|
||||
}
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
documentID := strconv.Itoa(hashSearXNGString(content, 100000000))
|
||||
displayID := strconv.Itoa(hashSearXNGString(documentID, 500))
|
||||
title := searxngText(item["title"])
|
||||
resultURL := searxngText(item["url"])
|
||||
chunks = append(chunks, map[string]any{
|
||||
"id": displayID,
|
||||
"chunk_id": documentID,
|
||||
"content": content,
|
||||
"doc_id": documentID,
|
||||
"docnm_kwd": title,
|
||||
"document_id": documentID,
|
||||
"document_name": title,
|
||||
"dataset_id": nil,
|
||||
"image_id": nil,
|
||||
"positions": nil,
|
||||
"url": resultURL,
|
||||
"similarity": 1,
|
||||
"vector_similarity": nil,
|
||||
"term_similarity": nil,
|
||||
"row_id": nil,
|
||||
"doc_type": nil,
|
||||
"document_metadata": nil,
|
||||
})
|
||||
docAggs = append(docAggs, map[string]any{
|
||||
"doc_name": title,
|
||||
"doc_id": documentID,
|
||||
"count": 1,
|
||||
"url": resultURL,
|
||||
})
|
||||
}
|
||||
return chunks, docAggs
|
||||
}
|
||||
|
||||
func renderSearXNGReferences(chunks []map[string]any, maxTokens int) string {
|
||||
if len(chunks) == 0 {
|
||||
return ""
|
||||
}
|
||||
usedTokens := 0
|
||||
blocks := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
content := searxngText(chunk["content"])
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
usedTokens += tokenizer.NumTokensFromString(content)
|
||||
var block strings.Builder
|
||||
fmt.Fprintf(&block, "\nID: %s", searxngText(chunk["id"]))
|
||||
if title := searxngPromptField(chunk["document_name"]); title != "" {
|
||||
fmt.Fprintf(&block, "\n├── Title: %s", title)
|
||||
}
|
||||
if resultURL := searxngPromptField(chunk["url"]); resultURL != "" {
|
||||
fmt.Fprintf(&block, "\n├── URL: %s", resultURL)
|
||||
}
|
||||
block.WriteString("\n└── Content:\n")
|
||||
block.WriteString(content)
|
||||
blocks = append(blocks, block.String())
|
||||
if maxTokens > 0 && float64(maxTokens)*0.97 < float64(usedTokens) {
|
||||
break
|
||||
}
|
||||
}
|
||||
return strings.Join(blocks, "\n")
|
||||
}
|
||||
|
||||
func searxngText(value any) string {
|
||||
if value == nil {
|
||||
return ""
|
||||
}
|
||||
if text, ok := value.(string); ok {
|
||||
return text
|
||||
}
|
||||
return fmt.Sprint(value)
|
||||
}
|
||||
|
||||
func searxngPromptField(value any) string {
|
||||
return searxngNewlinePattern.ReplaceAllString(searxngText(value), " ")
|
||||
}
|
||||
|
||||
func hashSearXNGString(value string, modulus int) int {
|
||||
digest := sha1.Sum([]byte(value))
|
||||
result := 0
|
||||
for _, part := range digest {
|
||||
result = (result*256 + int(part)) % modulus
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func init() {
|
||||
Register("SearXNG", newSearXNGComponent)
|
||||
}
|
||||
250
internal/agent/component/searxng_test.go
Normal file
250
internal/agent/component/searxng_test.go
Normal file
@@ -0,0 +1,250 @@
|
||||
//
|
||||
// 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 component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
|
||||
agentcanvas "ragflow/internal/agent/canvas"
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
type fakeSearXNGInvoker struct {
|
||||
args map[string]any
|
||||
calls int
|
||||
out string
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeSearXNGInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) {
|
||||
f.calls++
|
||||
if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return f.out, f.err
|
||||
}
|
||||
|
||||
func TestSearXNGRegisteredFactoryMatchesPythonSurface(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
component, err := New("SearXNG", map[string]any{
|
||||
"top_n": "10",
|
||||
"searxng_url": "http://localhost:4000",
|
||||
"outputs": map[string]any{"formalized_content": map[string]any{}},
|
||||
"setups": map[string]any{"query": "configured query"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New(SearXNG): %v", err)
|
||||
}
|
||||
if component.Name() != "SearXNG" {
|
||||
t.Fatalf("Name = %q, want SearXNG", component.Name())
|
||||
}
|
||||
if _, ok := component.Inputs()["query"]; !ok {
|
||||
t.Fatal("Inputs missing query")
|
||||
}
|
||||
if _, ok := component.Inputs()["searxng_url"]; !ok {
|
||||
t.Fatal("Inputs missing searxng_url")
|
||||
}
|
||||
for _, output := range []string{"formalized_content", "json"} {
|
||||
if _, ok := component.Outputs()[output]; !ok {
|
||||
t.Fatalf("Outputs missing %s", output)
|
||||
}
|
||||
}
|
||||
form := component.(*searxngComponent).GetInputForm()
|
||||
if len(form) != 2 {
|
||||
t.Fatalf("GetInputForm size = %d, want 2", len(form))
|
||||
}
|
||||
query := form["query"].(map[string]any)
|
||||
if query["name"] != "Query" || query["type"] != "line" {
|
||||
t.Fatalf("query form = %#v", query)
|
||||
}
|
||||
serverURL := form["searxng_url"].(map[string]any)
|
||||
if serverURL["name"] != "SearXNG URL" || serverURL["type"] != "line" || serverURL["placeholder"] != "http://localhost:4000" {
|
||||
t.Fatalf("searxng_url form = %#v", serverURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGInvokePreservesRawJSONPromptAndReferences(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
content := "RAGFlow content  remains"
|
||||
fake := &fakeSearXNGInvoker{out: `{"results":[{"title":"RAGFlow\nDocs","url":"https://ragflow.io","content":` + mustJSONText(t, content) + `,"engine":"bing","score":0.9},{"title":"Empty","url":"https://example.com","content":""}]}`}
|
||||
component := newSearXNGComponentWithInvoker(fake)
|
||||
state := runtime.NewCanvasState("run", "task")
|
||||
ctx := runtime.WithState(context.Background(), state)
|
||||
out, err := component.Invoke(ctx, map[string]any{
|
||||
"query": " ragflow ",
|
||||
"searxng_url": "http://localhost:4000",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if fake.args["query"] != " ragflow " || fake.args["searxng_url"] != "http://localhost:4000" {
|
||||
t.Fatalf("runtime args = %#v", fake.args)
|
||||
}
|
||||
results, ok := out["json"].([]any)
|
||||
if !ok || len(results) != 2 {
|
||||
t.Fatalf("json = %#v, want two raw results", out["json"])
|
||||
}
|
||||
first := results[0].(map[string]any)
|
||||
if first["engine"] != "bing" || first["score"] != float64(0.9) {
|
||||
t.Fatalf("raw result lost fields: %#v", first)
|
||||
}
|
||||
|
||||
cleaned := "RAGFlow content remains"
|
||||
documentID := hashSearXNGString(cleaned, 100000000)
|
||||
referenceID := hashSearXNGString(strconv.Itoa(documentID), 500)
|
||||
if documentID != 93760153 || referenceID != 491 {
|
||||
t.Fatalf("hash parity = %d/%d, want Python 93760153/491", documentID, referenceID)
|
||||
}
|
||||
formalized := out["formalized_content"].(string)
|
||||
for _, want := range []string{
|
||||
"ID: " + strconv.Itoa(referenceID),
|
||||
"Title: RAGFlow Docs",
|
||||
"URL: https://ragflow.io",
|
||||
"Content:\n" + cleaned,
|
||||
} {
|
||||
if !strings.Contains(formalized, want) {
|
||||
t.Fatalf("formalized_content missing %q: %s", want, formalized)
|
||||
}
|
||||
}
|
||||
|
||||
chunks := state.GetRetrievalChunks()
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("retrieval chunks = %#v, want one non-empty-content chunk", chunks)
|
||||
}
|
||||
if chunks[0]["document_id"] != strconv.Itoa(documentID) || chunks[0]["content"] != cleaned {
|
||||
t.Fatalf("retrieval chunk = %#v", chunks[0])
|
||||
}
|
||||
if chunks[0]["id"] != strconv.Itoa(referenceID) {
|
||||
t.Fatalf("chunk id = %#v, want displayed reference ID", chunks[0]["id"])
|
||||
}
|
||||
if chunks[0]["chunk_id"] != strconv.Itoa(documentID) || chunks[0]["doc_id"] != strconv.Itoa(documentID) {
|
||||
t.Fatalf("raw document IDs = %#v, want Python content hash", chunks[0])
|
||||
}
|
||||
if chunks[0]["similarity"] != 1 {
|
||||
t.Fatalf("similarity = %#v, want 1", chunks[0]["similarity"])
|
||||
}
|
||||
aggs := state.GetRetrievalDocAggs()
|
||||
if len(aggs) != 1 || aggs["RAGFlow\nDocs"]["doc_id"] != strconv.Itoa(documentID) {
|
||||
t.Fatalf("doc_aggs = %#v", aggs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGInvokeEmptyQuerySkipsTool(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fake := &fakeSearXNGInvoker{}
|
||||
component := newSearXNGComponentWithInvoker(fake)
|
||||
out, err := component.Invoke(context.Background(), map[string]any{"query": " "})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if fake.calls != 0 {
|
||||
t.Fatalf("tool calls = %d, want 0", fake.calls)
|
||||
}
|
||||
if out["formalized_content"] != "" {
|
||||
t.Fatalf("formalized_content = %#v", out["formalized_content"])
|
||||
}
|
||||
if results, ok := out["json"].([]any); !ok || len(results) != 0 {
|
||||
t.Fatalf("json = %#v, want empty []any", out["json"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGInvokePreservesErrorEnvelope(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fake := &fakeSearXNGInvoker{
|
||||
out: `{"results":[],"_ERROR":"Network error: upstream down"}`,
|
||||
err: errors.New("upstream down"),
|
||||
}
|
||||
component := newSearXNGComponentWithInvoker(fake)
|
||||
out, err := component.Invoke(context.Background(), map[string]any{"query": "ragflow"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if out["_ERROR"] != "Network error: upstream down" || out["formalized_content"] != "" {
|
||||
t.Fatalf("output = %#v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGPromptBoundaryMatchesPython(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
first := "small content"
|
||||
second := strings.Repeat("cross the token budget ", 50)
|
||||
third := "must not render"
|
||||
chunks := []map[string]any{
|
||||
{"id": "1", "content": first},
|
||||
{"id": "2", "content": second},
|
||||
{"id": "3", "content": third},
|
||||
}
|
||||
firstTokens := tokenizer.NumTokensFromString(first)
|
||||
maxTokens := int(math.Ceil(float64(firstTokens)/0.97)) + 1
|
||||
rendered := renderSearXNGReferences(chunks, maxTokens)
|
||||
if !strings.Contains(rendered, "ID: 1") || !strings.Contains(rendered, "ID: 2") {
|
||||
t.Fatalf("crossing chunk must be included like Python: %s", rendered)
|
||||
}
|
||||
if strings.Contains(rendered, "ID: 3") {
|
||||
t.Fatalf("chunk after budget crossing must be excluded: %s", rendered)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGBuildWorkflowUsesPythonComponentName(t *testing.T) {
|
||||
canvas := &agentcanvas.Canvas{
|
||||
Components: map[string]agentcanvas.CanvasComponent{
|
||||
"begin_0": {
|
||||
Obj: agentcanvas.CanvasComponentObj{ComponentName: "Begin", Params: map[string]any{}},
|
||||
Downstream: []string{"searxng_0"},
|
||||
},
|
||||
"searxng_0": {
|
||||
Obj: agentcanvas.CanvasComponentObj{ComponentName: "SearXNG", Params: map[string]any{
|
||||
"top_n": "10",
|
||||
"searxng_url": "http://localhost:4000",
|
||||
}},
|
||||
Upstream: []string{"begin_0"},
|
||||
Downstream: []string{"message_0"},
|
||||
},
|
||||
"message_0": {
|
||||
Obj: agentcanvas.CanvasComponentObj{ComponentName: "Message", Params: map[string]any{}},
|
||||
Upstream: []string{"searxng_0"},
|
||||
},
|
||||
},
|
||||
Path: []string{"begin_0", "searxng_0", "message_0"},
|
||||
}
|
||||
if _, err := agentcanvas.BuildWorkflow(context.Background(), canvas); err != nil {
|
||||
t.Fatalf("BuildWorkflow with SearXNG: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func mustJSONText(t *testing.T, value string) string {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
t.Fatalf("json.Marshal: %v", err)
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
109
internal/agent/component/wencai.go
Normal file
109
internal/agent/component/wencai.go
Normal file
@@ -0,0 +1,109 @@
|
||||
//
|
||||
// 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 component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
|
||||
agenttool "ragflow/internal/agent/tool"
|
||||
)
|
||||
|
||||
type wencaiInvoker interface {
|
||||
InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error)
|
||||
}
|
||||
|
||||
type wencaiComponent struct {
|
||||
inner wencaiInvoker
|
||||
}
|
||||
|
||||
func newWencaiComponent(params map[string]any) (Component, error) {
|
||||
toolParams := make(map[string]any, 2)
|
||||
for _, key := range []string{"top_n", "query_type"} {
|
||||
if value, ok := params[key]; ok {
|
||||
toolParams[key] = value
|
||||
}
|
||||
}
|
||||
inner, err := agenttool.BuildByName("wencai", toolParams)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
invoker, ok := inner.(wencaiInvoker)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("WenCai: tool does not implement InvokableRun")
|
||||
}
|
||||
return newWencaiComponentWithInvoker(invoker), nil
|
||||
}
|
||||
|
||||
func newWencaiComponentWithInvoker(inner wencaiInvoker) Component {
|
||||
return &wencaiComponent{inner: inner}
|
||||
}
|
||||
|
||||
func (c *wencaiComponent) Name() string { return "WenCai" }
|
||||
|
||||
func (c *wencaiComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"query": "The question/conditions to select stocks.",
|
||||
}
|
||||
}
|
||||
|
||||
func (c *wencaiComponent) Outputs() map[string]string {
|
||||
return map[string]string{
|
||||
"report": "WenCai query report.",
|
||||
}
|
||||
}
|
||||
|
||||
func (c *wencaiComponent) GetInputForm() map[string]any {
|
||||
return map[string]any{
|
||||
"query": map[string]any{
|
||||
"name": "Query",
|
||||
"type": "line",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *wencaiComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
query := stringParam(inputs["query"])
|
||||
if query == "" {
|
||||
return map[string]any{"report": ""}, nil
|
||||
}
|
||||
argsJSON, err := json.Marshal(map[string]any{"query": query})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("canvas: WenCai: encode query: %w", err)
|
||||
}
|
||||
out, err := c.inner.InvokableRun(ctx, string(argsJSON))
|
||||
decoded := parseToolEnvelope(out)
|
||||
report, _ := decoded["report"].(string)
|
||||
if message, _ := decoded["_ERROR"].(string); message != "" {
|
||||
return map[string]any{"report": report, "_ERROR": message}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("canvas: WenCai: %w", err)
|
||||
}
|
||||
return map[string]any{"report": report}, nil
|
||||
}
|
||||
|
||||
func (c *wencaiComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func init() {
|
||||
Register("WenCai", newWencaiComponent)
|
||||
}
|
||||
190
internal/agent/component/wencai_test.go
Normal file
190
internal/agent/component/wencai_test.go
Normal file
@@ -0,0 +1,190 @@
|
||||
//
|
||||
// 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 component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
|
||||
agentcanvas "ragflow/internal/agent/canvas"
|
||||
)
|
||||
|
||||
type fakeWencaiInvoker struct {
|
||||
args map[string]any
|
||||
calls int
|
||||
out string
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeWencaiInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) {
|
||||
f.calls++
|
||||
if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if f.out == "" && f.err == nil {
|
||||
return `{"report":""}`, nil
|
||||
}
|
||||
return f.out, f.err
|
||||
}
|
||||
|
||||
func TestWencai_RegisteredFactoryMatchesPythonSurface(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, err := New("WenCai", map[string]any{
|
||||
"top_n": float64(20),
|
||||
"query_type": "stock",
|
||||
"outputs": map[string]any{"report": map[string]any{}},
|
||||
"setups": map[string]any{"query": "configured query"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New(WenCai): %v", err)
|
||||
}
|
||||
if got := c.Name(); got != "WenCai" {
|
||||
t.Fatalf("Name() = %q, want WenCai", got)
|
||||
}
|
||||
if _, ok := c.Inputs()["query"]; !ok {
|
||||
t.Fatal("Inputs() missing query")
|
||||
}
|
||||
if _, ok := c.Outputs()["report"]; !ok {
|
||||
t.Fatal("Outputs() missing report")
|
||||
}
|
||||
formGetter, ok := c.(interface{ GetInputForm() map[string]any })
|
||||
if !ok {
|
||||
t.Fatal("WenCai component does not expose GetInputForm")
|
||||
}
|
||||
query, ok := formGetter.GetInputForm()["query"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("query form has type %T, want map", formGetter.GetInputForm()["query"])
|
||||
}
|
||||
if query["name"] != "Query" || query["type"] != "line" {
|
||||
t.Fatalf("query form = %#v, want name=Query type=line", query)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_InvokeReturnsEmptyReportWithoutError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fake := &fakeWencaiInvoker{}
|
||||
c := newWencaiComponentWithInvoker(fake)
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"query": "商业航天",
|
||||
"top_n": 20,
|
||||
"query_type": "stock",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got := out["report"]; got != "" {
|
||||
t.Fatalf("report = %v, want empty string", got)
|
||||
}
|
||||
if fake.calls != 1 {
|
||||
t.Fatalf("calls = %d, want 1", fake.calls)
|
||||
}
|
||||
if got := fake.args["query"]; got != "商业航天" {
|
||||
t.Fatalf("runtime query = %v, want 商业航天", got)
|
||||
}
|
||||
if len(fake.args) != 1 {
|
||||
t.Fatalf("runtime args = %#v, want query only", fake.args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_InvokeEmptyQuerySkipsTool(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fake := &fakeWencaiInvoker{}
|
||||
c := newWencaiComponentWithInvoker(fake)
|
||||
out, err := c.Invoke(context.Background(), map[string]any{"query": ""})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got := out["report"]; got != "" {
|
||||
t.Fatalf("report = %v, want empty string", got)
|
||||
}
|
||||
if fake.calls != 0 {
|
||||
t.Fatalf("calls = %d, want 0", fake.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_InvokePreservesErrorEnvelope(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fake := &fakeWencaiInvoker{
|
||||
out: `{"report":"","_ERROR":"upstream failed"}`,
|
||||
err: errors.New("upstream failed"),
|
||||
}
|
||||
c := newWencaiComponentWithInvoker(fake)
|
||||
out, err := c.Invoke(context.Background(), map[string]any{"query": "商业航天"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got := out["report"]; got != "" {
|
||||
t.Fatalf("report = %v, want empty string", got)
|
||||
}
|
||||
if got := out["_ERROR"]; got != "upstream failed" {
|
||||
t.Fatalf("_ERROR = %v, want upstream failed", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_InvokePreservesErrorEnvelopeWithoutGoError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fake := &fakeWencaiInvoker{
|
||||
out: `{"report":"","_ERROR":"business error"}`,
|
||||
}
|
||||
c := newWencaiComponentWithInvoker(fake)
|
||||
out, err := c.Invoke(context.Background(), map[string]any{"query": "商业航天"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got := out["report"]; got != "" {
|
||||
t.Fatalf("report = %v, want empty string", got)
|
||||
}
|
||||
if got := out["_ERROR"]; got != "business error" {
|
||||
t.Fatalf("_ERROR = %v, want business error", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_BuildWorkflowUsesPythonComponentName(t *testing.T) {
|
||||
c := &agentcanvas.Canvas{
|
||||
Components: map[string]agentcanvas.CanvasComponent{
|
||||
"begin_0": {
|
||||
Obj: agentcanvas.CanvasComponentObj{ComponentName: "Begin", Params: map[string]any{}},
|
||||
Downstream: []string{"wencai_0"},
|
||||
},
|
||||
"wencai_0": {
|
||||
Obj: agentcanvas.CanvasComponentObj{ComponentName: "WenCai", Params: map[string]any{
|
||||
"top_n": float64(20),
|
||||
"query_type": "stock",
|
||||
}},
|
||||
Upstream: []string{"begin_0"},
|
||||
Downstream: []string{"message_0"},
|
||||
},
|
||||
"message_0": {
|
||||
Obj: agentcanvas.CanvasComponentObj{ComponentName: "Message", Params: map[string]any{}},
|
||||
Upstream: []string{"wencai_0"},
|
||||
},
|
||||
},
|
||||
Path: []string{"begin_0", "wencai_0", "message_0"},
|
||||
}
|
||||
if _, err := agentcanvas.BuildWorkflow(context.Background(), c); err != nil {
|
||||
t.Fatalf("BuildWorkflow with WenCai: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -405,9 +405,26 @@ func (s *CanvasState) SetRetrievalReferences(chunks, docAggs []map[string]any) {
|
||||
if chunkValues == nil {
|
||||
chunkValues = make([]any, 0, len(chunks))
|
||||
}
|
||||
seenChunkIDs := make(map[string]struct{}, len(chunkValues)+len(chunks))
|
||||
for _, value := range chunkValues {
|
||||
chunk, ok := value.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if id, ok := retrievalReferenceID(chunk); ok {
|
||||
seenChunkIDs[id] = struct{}{}
|
||||
}
|
||||
}
|
||||
for _, chunk := range chunks {
|
||||
if id, ok := retrievalReferenceID(chunk); ok {
|
||||
if _, exists := seenChunkIDs[id]; exists {
|
||||
continue
|
||||
}
|
||||
seenChunkIDs[id] = struct{}{}
|
||||
}
|
||||
chunkValues = append(chunkValues, chunk)
|
||||
}
|
||||
|
||||
docAggValues, _ := s.Retrieval["doc_aggs"].(map[string]any)
|
||||
if docAggValues == nil {
|
||||
docAggValues = make(map[string]any, len(docAggs))
|
||||
@@ -427,6 +444,35 @@ func (s *CanvasState) SetRetrievalReferences(chunks, docAggs []map[string]any) {
|
||||
s.Retrieval["doc_aggs"] = docAggValues
|
||||
}
|
||||
|
||||
func retrievalReferenceID(chunk map[string]any) (string, bool) {
|
||||
value, ok := chunk["id"]
|
||||
if !ok || value == nil {
|
||||
return "", false
|
||||
}
|
||||
id := fmt.Sprint(value)
|
||||
return id, id != ""
|
||||
}
|
||||
|
||||
// GetRetrievalDocAggs returns a shallow snapshot keyed by document name.
|
||||
func (s *CanvasState) GetRetrievalDocAggs() map[string]map[string]any {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
raw, _ := s.Retrieval["doc_aggs"].(map[string]any)
|
||||
if raw == nil {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]map[string]any, len(raw))
|
||||
for name, item := range raw {
|
||||
if agg, ok := item.(map[string]any); ok {
|
||||
out[name] = agg
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// getVarLocked is the lock-free inner GetVar. Caller must hold s.mu (read or
|
||||
// write) for the entire call.
|
||||
func getVarLocked(s *CanvasState, ref string) (any, error) {
|
||||
|
||||
@@ -94,41 +94,43 @@ func TestCanvasState_SetRetrievalReferencesMergesCalls(t *testing.T) {
|
||||
[]map[string]any{{"doc_name": "doc-1", "count": 1}},
|
||||
)
|
||||
state.SetRetrievalReferences(
|
||||
[]map[string]any{{"id": "chunk-2"}},
|
||||
[]map[string]any{
|
||||
{"id": "chunk-1", "content": "duplicate"},
|
||||
{"id": "chunk-2"},
|
||||
{"content": "chunk without an ID"},
|
||||
},
|
||||
[]map[string]any{
|
||||
{"doc_name": "doc-1", "count": 99},
|
||||
{"doc_name": "doc-2", "count": 2},
|
||||
{"doc_name": "", "count": 3},
|
||||
},
|
||||
)
|
||||
|
||||
chunks := state.GetRetrievalChunks()
|
||||
if len(chunks) != 2 {
|
||||
t.Fatalf("chunks length = %d, want 2", len(chunks))
|
||||
if len(chunks) != 3 {
|
||||
t.Fatalf("chunks length = %d, want 3", len(chunks))
|
||||
}
|
||||
if chunks[0]["id"] != "chunk-1" || chunks[1]["id"] != "chunk-2" {
|
||||
t.Errorf("chunk IDs = [%v %v], want [chunk-1 chunk-2]", chunks[0]["id"], chunks[1]["id"])
|
||||
}
|
||||
|
||||
state.mu.RLock()
|
||||
rawDocAggs := state.Retrieval["doc_aggs"]
|
||||
docAggs, ok := rawDocAggs.(map[string]any)
|
||||
state.mu.RUnlock()
|
||||
if !ok {
|
||||
t.Fatalf("doc_aggs type = %T, want map[string]any", rawDocAggs)
|
||||
if chunks[2]["content"] != "chunk without an ID" {
|
||||
t.Errorf("chunk without ID was not retained: %#v", chunks[2])
|
||||
}
|
||||
|
||||
docAggs := state.GetRetrievalDocAggs()
|
||||
if len(docAggs) != 2 {
|
||||
t.Fatalf("doc_aggs length = %d, want 2", len(docAggs))
|
||||
}
|
||||
firstDoc, ok := docAggs["doc-1"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("doc_aggs[doc-1] type = %T, want map[string]any", docAggs["doc-1"])
|
||||
}
|
||||
firstDoc := docAggs["doc-1"]
|
||||
if firstDoc["count"] != 1 {
|
||||
t.Errorf("doc_aggs[doc-1].count = %v, want first value 1", firstDoc["count"])
|
||||
}
|
||||
if _, ok := docAggs["doc-2"]; !ok {
|
||||
t.Error("doc_aggs missing doc-2 from the second call")
|
||||
}
|
||||
if _, ok := docAggs[""]; ok {
|
||||
t.Error("doc_aggs retained an empty document name")
|
||||
}
|
||||
}
|
||||
|
||||
func contains(s, substr string) bool {
|
||||
|
||||
@@ -51,11 +51,11 @@ var registry = map[string]Factory{
|
||||
"retrieval": noConfig("retrieval", func() einotool.BaseTool { return NewRetrievalTool() }),
|
||||
"search_my_dataset": noConfig("search_my_dataset", 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() }),
|
||||
"searxng": buildSearXNGTool,
|
||||
"tavily": noConfig("tavily", func() einotool.BaseTool { return NewTavilyTool() }),
|
||||
"tavily_extract": noConfig("tavily_extract", func() einotool.BaseTool { return NewTavilyExtractTool() }),
|
||||
"tushare": noConfig("tushare", func() einotool.BaseTool { return NewTushareTool() }),
|
||||
"wencai": noConfig("wencai", func() einotool.BaseTool { return NewWencaiTool() }),
|
||||
"wencai": buildWencaiTool,
|
||||
"web_crawler": noConfig("web_crawler", func() einotool.BaseTool { return NewCrawlerTool() }),
|
||||
"wikipedia": buildWikipediaTool,
|
||||
"wikipedia_search": buildWikipediaTool,
|
||||
@@ -278,6 +278,58 @@ func buildPubMedTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
return NewPubMedToolWithDefaults(nil, defaults), nil
|
||||
}
|
||||
|
||||
func buildSearXNGTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
defaults := defaultSearXNGParams()
|
||||
for key := range params {
|
||||
switch key {
|
||||
case "top_n", "searxng_url":
|
||||
default:
|
||||
return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "searxng", key)
|
||||
}
|
||||
}
|
||||
if value, ok := params["top_n"]; ok {
|
||||
topN, valid := parseSearXNGTopN(value)
|
||||
if !valid || topN <= 0 {
|
||||
return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param top_n", "searxng")
|
||||
}
|
||||
defaults.TopN = topN
|
||||
}
|
||||
if value, ok := params["searxng_url"]; ok {
|
||||
searxngURL, valid := value.(string)
|
||||
if !valid {
|
||||
return nil, fmt.Errorf("agent tool: tool %q requires string node-level param searxng_url", "searxng")
|
||||
}
|
||||
defaults.SearXNGURL = searxngURL
|
||||
}
|
||||
return newSearXNGToolWithDefaults(nil, defaults), nil
|
||||
}
|
||||
|
||||
func buildWencaiTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
defaults := wencaiParams{}
|
||||
for key := range params {
|
||||
switch key {
|
||||
case "top_n", "query_type":
|
||||
default:
|
||||
return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "wencai", key)
|
||||
}
|
||||
}
|
||||
if value, ok := params["top_n"]; ok {
|
||||
topN, valid := strictInt(value)
|
||||
if !valid || topN <= 0 {
|
||||
return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param top_n", "wencai")
|
||||
}
|
||||
defaults.TopN = topN
|
||||
}
|
||||
if value, ok := params["query_type"]; ok {
|
||||
queryType, valid := value.(string)
|
||||
if !valid || !isWencaiQueryTypeSupported(queryType) {
|
||||
return nil, fmt.Errorf("agent tool: tool %q has unsupported query_type %q", "wencai", queryType)
|
||||
}
|
||||
defaults.QueryType = queryType
|
||||
}
|
||||
return newWencaiTool(defaults), nil
|
||||
}
|
||||
|
||||
func buildKeenableTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
if len(params) == 0 {
|
||||
return NewKeenableTool(), nil
|
||||
@@ -383,6 +435,29 @@ func intParam(params map[string]any, key string) (int, bool) {
|
||||
}
|
||||
}
|
||||
|
||||
func strictInt(value any) (int, bool) {
|
||||
switch x := value.(type) {
|
||||
case int:
|
||||
return x, true
|
||||
case int32:
|
||||
return int(x), true
|
||||
case int64:
|
||||
if int64(int(x)) != x {
|
||||
return 0, false
|
||||
}
|
||||
return int(x), true
|
||||
case float64:
|
||||
maxInt := int(^uint(0) >> 1)
|
||||
minInt := -maxInt - 1
|
||||
if math.Trunc(x) != x || x >= float64(maxInt) || x <= float64(minInt) {
|
||||
return 0, false
|
||||
}
|
||||
return int(x), true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
func boolParam(params map[string]any, key string) (bool, bool) {
|
||||
v, ok := params[key]
|
||||
if !ok {
|
||||
|
||||
@@ -20,67 +20,77 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const searxngToolName = "searxng"
|
||||
const (
|
||||
searxngToolName = "searxng_search"
|
||||
searxngToolDescription = "SearXNG is a privacy-focused metasearch engine that aggregates results from multiple search engines without tracking users. It provides comprehensive web search capabilities."
|
||||
defaultSearXNGTopN = 10
|
||||
searxngRequestTimeout = 10 * time.Second
|
||||
)
|
||||
|
||||
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.
|
||||
// searxngParams mirrors the SearXNG-specific Python parameters. Query and
|
||||
// searxng_url are model inputs; searxng_url may also be node configuration,
|
||||
// while top_n is Canvas-only configuration.
|
||||
type searxngParams struct {
|
||||
BaseURL string `json:"base_url"`
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
SearXNGURL string `json:"searxng_url"`
|
||||
TopN int `json:"top_n"`
|
||||
}
|
||||
|
||||
// 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"`
|
||||
Results []any `json:"results"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// SearXNGTool is the SearXNG
|
||||
// meta-search tool. It calls
|
||||
// a self-hosted SearXNG instance via the shared HTTPHelper.
|
||||
type searxngResolver func(string) (string, net.IP, error)
|
||||
|
||||
// SearXNGTool queries a caller-configured SearXNG instance.
|
||||
type SearXNGTool struct {
|
||||
helper *HTTPHelper
|
||||
helper *HTTPHelper
|
||||
defaults searxngParams
|
||||
resolve searxngResolver
|
||||
}
|
||||
|
||||
// NewSearXNGTool returns a SearXNGTool using the default HTTPHelper.
|
||||
func defaultSearXNGParams() searxngParams {
|
||||
return searxngParams{TopN: defaultSearXNGTopN}
|
||||
}
|
||||
|
||||
// NewSearXNGTool returns a SearXNG tool using Python's defaults.
|
||||
func NewSearXNGTool() *SearXNGTool {
|
||||
return NewSearXNGToolWith(NewHTTPHelper())
|
||||
return newSearXNGToolWithDefaults(nil, defaultSearXNGParams())
|
||||
}
|
||||
|
||||
// NewSearXNGToolWith returns a SearXNGTool that uses the provided
|
||||
// HTTPHelper. Useful for tests.
|
||||
func NewSearXNGToolWith(h *HTTPHelper) *SearXNGTool {
|
||||
if h == nil {
|
||||
h = NewHTTPHelper()
|
||||
func newSearXNGToolWithDefaults(helper *HTTPHelper, defaults searxngParams) *SearXNGTool {
|
||||
if helper == nil {
|
||||
helper = NewHTTPHelper()
|
||||
}
|
||||
if defaults.TopN == 0 {
|
||||
defaults.TopN = defaultSearXNGTopN
|
||||
}
|
||||
// SearXNG performs one request. Keep the shared helper from adding its own
|
||||
// retry loop so request count stays explicit and predictable.
|
||||
helper.retry.MaxAttempts = 1
|
||||
if helper.client != nil {
|
||||
helper.client.Timeout = searxngRequestTimeout
|
||||
}
|
||||
return &SearXNGTool{
|
||||
helper: helper,
|
||||
defaults: defaults,
|
||||
resolve: ResolveAndValidate,
|
||||
}
|
||||
return &SearXNGTool{helper: h}
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
// Info exposes the Python model-call schema, not Canvas-only configuration.
|
||||
func (s *SearXNGTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: searxngToolName,
|
||||
@@ -88,81 +98,131 @@ func (s *SearXNGTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query",
|
||||
Desc: "The search keywords to execute with SearXNG. The keywords should be the most important words/terms(includes synonyms) from the original request.",
|
||||
Required: true,
|
||||
},
|
||||
"base_url": {
|
||||
"searxng_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.",
|
||||
Desc: "The base URL of your SearXNG instance (e.g., http://localhost:4000). This is required to connect to your SearXNG server.",
|
||||
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()
|
||||
func buildSearXNGURL(baseURL, query string) string {
|
||||
params := url.Values{}
|
||||
params.Set("q", query)
|
||||
params.Set("format", "json")
|
||||
params.Set("categories", "general")
|
||||
params.Set("language", "auto")
|
||||
params.Set("safesearch", "1")
|
||||
params.Set("pageno", "1")
|
||||
return baseURL + "/search?" + params.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)
|
||||
func mergeSearXNGDefaults(defaults, params searxngParams) searxngParams {
|
||||
if defaults.SearXNGURL != "" {
|
||||
params.SearXNGURL = defaults.SearXNGURL
|
||||
}
|
||||
if p.Query == "" {
|
||||
return searxngErrJSON(fmt.Errorf("query is required")),
|
||||
fmt.Errorf("searxng: query is required")
|
||||
params.TopN = defaults.TopN
|
||||
return params
|
||||
}
|
||||
|
||||
// InvokableRun performs the same request and result slicing as Python
|
||||
// SearXNG._invoke. Empty try-run inputs return an empty result without I/O.
|
||||
func (s *SearXNGTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
var params searxngParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), ¶ms); err != nil {
|
||||
err = fmt.Errorf("searxng: parse arguments: %w", err)
|
||||
return searxngErrJSON(err), err
|
||||
}
|
||||
params = mergeSearXNGDefaults(s.defaults, params)
|
||||
if strings.TrimSpace(params.Query) == "" {
|
||||
return searxngJSON(searxngEnvelope{Results: []any{}}), nil
|
||||
}
|
||||
params.SearXNGURL = strings.TrimSpace(params.SearXNGURL)
|
||||
if params.SearXNGURL == "" {
|
||||
return searxngJSON(searxngEnvelope{Results: []any{}}), nil
|
||||
}
|
||||
|
||||
endpoint := buildSearXNGURL(p.BaseURL, p.Query, p.MaxResults)
|
||||
resp, err := s.helper.Do(ctx, http.MethodGet, endpoint, "", "", nil)
|
||||
host, pinnedIP, err := s.resolve(params.SearXNGURL)
|
||||
if err != nil {
|
||||
return searxngErrJSON(err), err
|
||||
}
|
||||
|
||||
endpoint := buildSearXNGURL(params.SearXNGURL, params.Query)
|
||||
if ctx.Err() != nil {
|
||||
return searxngJSON(searxngEnvelope{Results: []any{}}), nil
|
||||
}
|
||||
results, requestErr := s.fetch(ctx, endpoint, host, pinnedIP)
|
||||
if requestErr != nil {
|
||||
if ctx.Err() != nil {
|
||||
return searxngJSON(searxngEnvelope{Results: []any{}}), nil
|
||||
}
|
||||
return searxngErrJSON(requestErr), requestErr
|
||||
}
|
||||
if len(results) > params.TopN {
|
||||
results = results[:params.TopN]
|
||||
}
|
||||
return searxngJSON(searxngEnvelope{Results: results}), nil
|
||||
}
|
||||
|
||||
func (s *SearXNGTool) fetch(ctx context.Context, endpoint, host string, pinnedIP net.IP) ([]any, error) {
|
||||
requestCtx, cancel := context.WithTimeout(ctx, searxngRequestTimeout)
|
||||
defer cancel()
|
||||
|
||||
resp, err := s.helper.DoPinned(requestCtx, http.MethodGet, endpoint, "", "", nil, host, pinnedIP)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Network error: %w", 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)
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
return nil, fmt.Errorf("Network error: SearXNG 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)
|
||||
var data map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
|
||||
return nil, fmt.Errorf("Invalid response from SearXNG: %w", err)
|
||||
}
|
||||
return searxngJSON(searxngEnvelope{Results: raw.Results}), nil
|
||||
if data == nil {
|
||||
return nil, fmt.Errorf("Invalid response from SearXNG")
|
||||
}
|
||||
rawResults, ok := data["results"]
|
||||
if !ok {
|
||||
return []any{}, nil
|
||||
}
|
||||
results, ok := rawResults.([]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("Invalid results format from SearXNG")
|
||||
}
|
||||
for _, result := range results {
|
||||
if _, ok := result.(map[string]any); !ok {
|
||||
return nil, fmt.Errorf("Invalid results format from SearXNG")
|
||||
}
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func parseSearXNGTopN(value any) (int, bool) {
|
||||
if text, ok := value.(string); ok {
|
||||
parsed, err := strconv.Atoi(strings.TrimSpace(text))
|
||||
return parsed, err == nil
|
||||
}
|
||||
return strictInt(value)
|
||||
}
|
||||
|
||||
func searxngJSON(env searxngEnvelope) string {
|
||||
b, err := json.Marshal(env)
|
||||
data, err := json.Marshal(env)
|
||||
if err != nil {
|
||||
return fmt.Sprintf(`{"_ERROR":"searxng: marshal result: %s"}`, err)
|
||||
}
|
||||
return string(b)
|
||||
return string(data)
|
||||
}
|
||||
|
||||
func searxngErrJSON(err error) string {
|
||||
return searxngJSON(searxngEnvelope{Error: err.Error()})
|
||||
if err == nil {
|
||||
err = fmt.Errorf("SearXNG error")
|
||||
}
|
||||
return searxngJSON(searxngEnvelope{Results: []any{}, Error: err.Error()})
|
||||
}
|
||||
|
||||
@@ -19,138 +19,233 @@ package tool
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSearXNG_BuildURL(t *testing.T) {
|
||||
func TestSearXNGBuildURLMatchesPythonQuery(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",
|
||||
},
|
||||
got := buildSearXNGURL("https://searx.example.com", "rag flow")
|
||||
parsed, err := url.Parse(got)
|
||||
if err != nil {
|
||||
t.Fatalf("url.Parse(%q): %v", got, err)
|
||||
}
|
||||
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"))
|
||||
}
|
||||
})
|
||||
if parsed.Path != "/search" {
|
||||
t.Fatalf("path = %q, want /search", parsed.Path)
|
||||
}
|
||||
want := map[string]string{
|
||||
"q": "rag flow",
|
||||
"format": "json",
|
||||
"categories": "general",
|
||||
"language": "auto",
|
||||
"safesearch": "1",
|
||||
"pageno": "1",
|
||||
}
|
||||
for key, value := range want {
|
||||
if actual := parsed.Query().Get(key); actual != value {
|
||||
t.Errorf("query[%s] = %q, want %q", key, actual, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNG_ParseResults(t *testing.T) {
|
||||
func TestSearXNGInvokableRunPreservesRawResultsAndTopN(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
if request.URL.Path != "/search" {
|
||||
t.Errorf("path = %q, want /search", request.URL.Path)
|
||||
}
|
||||
if request.URL.Query().Get("q") != " ragflow search " {
|
||||
t.Errorf("q = %q, want original query whitespace", request.URL.Query().Get("q"))
|
||||
}
|
||||
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"}
|
||||
{"title":"RAGFlow","url":"https://ragflow.io","content":"Open source RAG engine","engine":"bing","score":0.9},
|
||||
{"title":"GitHub","url":"https://github.com/infiniflow/ragflow","content":"Source code","engines":["google","bing"]},
|
||||
{"title":"Dropped","url":"https://example.com","content":"top_n applies"}
|
||||
]
|
||||
}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer server.Close()
|
||||
|
||||
tool := NewSearXNGTool()
|
||||
out, err := tool.InvokableRun(context.Background(),
|
||||
`{"query":"ragflow","base_url":`+jsonString(srv.URL)+`,"max_results":5}`)
|
||||
defaults := defaultSearXNGParams()
|
||||
defaults.SearXNGURL = server.URL
|
||||
defaults.TopN = 2
|
||||
searchTool := newLocalSearXNGTool(t, defaults)
|
||||
out, err := searchTool.InvokableRun(context.Background(), `{"query":" ragflow search ","searxng_url":"http://127.0.0.1:1"}`)
|
||||
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)
|
||||
var envelope searxngEnvelope
|
||||
if err := json.Unmarshal([]byte(out), &envelope); err != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", err, out)
|
||||
}
|
||||
if env.Error != "" {
|
||||
t.Errorf("Error = %q, want empty", env.Error)
|
||||
if len(envelope.Results) != 2 {
|
||||
t.Fatalf("results len = %d, want 2", len(envelope.Results))
|
||||
}
|
||||
if len(env.Results) != 2 {
|
||||
t.Fatalf("Results len = %d, want 2", len(env.Results))
|
||||
first := envelope.Results[0].(map[string]any)
|
||||
if first["engine"] != "bing" || first["score"] != float64(0.9) {
|
||||
t.Fatalf("first result lost upstream fields: %#v", first)
|
||||
}
|
||||
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)
|
||||
second := envelope.Results[1].(map[string]any)
|
||||
if _, ok := second["engines"].([]any); !ok {
|
||||
t.Fatalf("second result lost engines: %#v", second)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNG_Info(t *testing.T) {
|
||||
func TestSearXNGInfoMatchesPythonModelSchema(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewSearXNGTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
info, err := NewSearXNGTool().Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "searxng" {
|
||||
t.Errorf("Name = %q, want searxng", info.Name)
|
||||
if info.Name != "searxng_search" {
|
||||
t.Fatalf("Name = %q, want searxng_search", info.Name)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "SearXNG") {
|
||||
t.Errorf("Desc = %q, want to mention SearXNG", info.Desc)
|
||||
if info.Desc != searxngToolDescription {
|
||||
t.Fatalf("Desc = %q, want Python description", info.Desc)
|
||||
}
|
||||
jsonSchema, err := info.ParamsOneOf.ToJSONSchema()
|
||||
if err != nil {
|
||||
t.Fatalf("ToJSONSchema: %v", err)
|
||||
}
|
||||
raw, err := json.Marshal(jsonSchema)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal schema: %v", err)
|
||||
}
|
||||
schemaText := string(raw)
|
||||
for _, key := range []string{`"query"`, `"searxng_url"`, `"required":["query"]`} {
|
||||
if !strings.Contains(schemaText, key) {
|
||||
t.Fatalf("schema missing %s: %s", key, schemaText)
|
||||
}
|
||||
}
|
||||
if strings.Contains(schemaText, `"top_n"`) {
|
||||
t.Fatalf("schema leaked node config top_n: %s", schemaText)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNG_RequiresQuery(t *testing.T) {
|
||||
func TestSearXNGEmptyTryRunInputsSkipRequest(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewSearXNGTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{"query":""}`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty query")
|
||||
searchTool := NewSearXNGTool()
|
||||
var resolves atomic.Int32
|
||||
searchTool.resolve = func(string) (string, net.IP, error) {
|
||||
resolves.Add(1)
|
||||
return "", nil, errors.New("must not resolve")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "query") {
|
||||
t.Errorf("err = %v, want to mention query", err)
|
||||
for _, args := range []string{
|
||||
`{"query":""}`,
|
||||
`{"query":" ","searxng_url":"https://example.com"}`,
|
||||
`{"query":"ragflow"}`,
|
||||
} {
|
||||
out, err := searchTool.InvokableRun(context.Background(), args)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun(%s): %v", args, err)
|
||||
}
|
||||
var envelope searxngEnvelope
|
||||
if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 || envelope.Error != "" {
|
||||
t.Fatalf("InvokableRun(%s) = %s, want empty results", args, out)
|
||||
}
|
||||
}
|
||||
if resolves.Load() != 0 {
|
||||
t.Fatalf("resolver calls = %d, want 0", resolves.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGBuildByNameAcceptsPythonNodeParams(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
built, err := BuildByName("searxng", map[string]any{
|
||||
"top_n": "8",
|
||||
"searxng_url": "https://searx.example.com",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildByName: %v", err)
|
||||
}
|
||||
searchTool, ok := built.(*SearXNGTool)
|
||||
if !ok {
|
||||
t.Fatalf("built type = %T, want *SearXNGTool", built)
|
||||
}
|
||||
if searchTool.defaults.TopN != 8 || searchTool.defaults.SearXNGURL != "https://searx.example.com" {
|
||||
t.Fatalf("defaults = %+v", searchTool.defaults)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGBuildByNameRejectsInvalidNodeParams(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
invalid := []map[string]any{
|
||||
{"top_n": "abc"},
|
||||
{"top_n": 0},
|
||||
{"top_n": 1.5},
|
||||
{"unknown": true},
|
||||
}
|
||||
for _, params := range invalid {
|
||||
if _, err := BuildByName("searxng", params); err == nil {
|
||||
t.Fatalf("BuildByName(%#v) succeeded, want validation error", params)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGDoesNotRetryFailedRequest(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var calls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
calls.Add(1)
|
||||
http.Error(w, "temporary", http.StatusServiceUnavailable)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
defaults := defaultSearXNGParams()
|
||||
defaults.SearXNGURL = server.URL
|
||||
searchTool := newLocalSearXNGTool(t, defaults)
|
||||
if _, err := searchTool.InvokableRun(context.Background(), `{"query":"single attempt"}`); err == nil {
|
||||
t.Fatal("InvokableRun succeeded, want upstream error")
|
||||
}
|
||||
if calls.Load() != 1 {
|
||||
t.Fatalf("request calls = %d, want 1", calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearXNGSSRFGuardRejectsLoopback(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
defaults := defaultSearXNGParams()
|
||||
defaults.SearXNGURL = "http://127.0.0.1:4000"
|
||||
searchTool := newSearXNGToolWithDefaults(nil, defaults)
|
||||
out, err := searchTool.InvokableRun(context.Background(), `{"query":"metadata"}`)
|
||||
if err == nil || !errors.Is(err, ErrSSRFBlocked) {
|
||||
t.Fatalf("err = %v, want ErrSSRFBlocked", err)
|
||||
}
|
||||
if !strings.Contains(out, `"_ERROR"`) {
|
||||
t.Fatalf("output = %s, want _ERROR envelope", out)
|
||||
}
|
||||
}
|
||||
|
||||
func newLocalSearXNGTool(t *testing.T, defaults searxngParams) *SearXNGTool {
|
||||
t.Helper()
|
||||
searchTool := newSearXNGToolWithDefaults(nil, defaults)
|
||||
searchTool.resolve = func(rawURL string) (string, net.IP, error) {
|
||||
parsed, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
ip := net.ParseIP(parsed.Hostname())
|
||||
if ip == nil {
|
||||
return "", nil, errors.New("test URL must use a literal IP")
|
||||
}
|
||||
return parsed.Hostname(), ip, nil
|
||||
}
|
||||
return searchTool
|
||||
}
|
||||
|
||||
@@ -19,100 +19,129 @@ package tool
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const wencaiToolName = "wencai"
|
||||
const (
|
||||
wencaiToolName = "iwencai"
|
||||
defaultWencaiTopN = 10
|
||||
defaultWencaiQueryType = "stock"
|
||||
)
|
||||
|
||||
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 wencaiToolDescription = `
|
||||
iwencai search: search platform is committed to providing hundreds of millions of investors with the most timely, accurate and comprehensive information, covering news, announcements, research reports, blogs, forums, Weibo, characters, etc.
|
||||
robo-advisor intelligent stock selection platform: through AI technology, is committed to providing investors with intelligent stock selection, quantitative investment, main force tracking, value investment, technical analysis and other types of stock selection technologies.
|
||||
fund selection platform: through AI technology, is committed to providing excellent fund, value investment, quantitative analysis and other fund selection technologies for foundation citizens.
|
||||
`
|
||||
|
||||
const wencaiUnsupportedMessage = "Wencai requires web scraping of 同花顺 — not yet implemented in Go Canvas. " +
|
||||
"Use Python Canvas."
|
||||
var wencaiQueryTypes = map[string]struct{}{
|
||||
"stock": {},
|
||||
"zhishu": {},
|
||||
"fund": {},
|
||||
"hkstock": {},
|
||||
"usstock": {},
|
||||
"threeboard": {},
|
||||
"conbond": {},
|
||||
"insurance": {},
|
||||
"futures": {},
|
||||
"lccp": {},
|
||||
"foreign_exchange": {},
|
||||
}
|
||||
|
||||
// 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.
|
||||
// wencaiParams mirrors Python WenCaiParam. Query is model-provided runtime
|
||||
// input; top_n and query_type are Canvas node configuration.
|
||||
type wencaiParams struct {
|
||||
Query string `json:"query"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PerPage int `json:"per_page,omitempty"`
|
||||
Query string `json:"query"`
|
||||
TopN int `json:"top_n"`
|
||||
QueryType string `json:"query_type"`
|
||||
}
|
||||
|
||||
// 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"`
|
||||
Report string `json:"report"`
|
||||
Error string `json:"_ERROR,omitempty"`
|
||||
}
|
||||
|
||||
// WencaiTool is a stub for the
|
||||
// 同花顺 Wencai (问财) natural-language stock screening tool
|
||||
// ( .
|
||||
//
|
||||
// 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{}
|
||||
// WencaiTool implements the behavior currently exposed by Python WenCai.
|
||||
// The Python integration has its upstream request disabled and returns an
|
||||
// empty report for both empty and non-empty queries, so the Go tool does the
|
||||
// same without reporting a false unsupported error.
|
||||
type WencaiTool struct {
|
||||
defaults wencaiParams
|
||||
}
|
||||
|
||||
// NewWencaiTool returns a WencaiTool. No HTTPHelper is allocated; the
|
||||
// stub never issues network requests.
|
||||
func NewWencaiTool() *WencaiTool { return &WencaiTool{} }
|
||||
func NewWencaiTool() *WencaiTool {
|
||||
return newWencaiTool(wencaiParams{})
|
||||
}
|
||||
|
||||
// Info returns the tool's metadata for the chat model.
|
||||
func newWencaiTool(defaults wencaiParams) *WencaiTool {
|
||||
if defaults.TopN == 0 {
|
||||
defaults.TopN = defaultWencaiTopN
|
||||
}
|
||||
if strings.TrimSpace(defaults.QueryType) == "" {
|
||||
defaults.QueryType = defaultWencaiQueryType
|
||||
}
|
||||
return &WencaiTool{defaults: defaults}
|
||||
}
|
||||
|
||||
// Info exposes only Python meta.parameters. Node configuration does not
|
||||
// belong in the model-emitted function-call schema.
|
||||
func (w *WencaiTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: wencaiToolName,
|
||||
Desc: wencaiToolDescription,
|
||||
Desc: strings.TrimSpace(wencaiToolDescription),
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Natural-language Wencai query (e.g. \"近期涨停股\", \"高股息低估值\").",
|
||||
Desc: "The question/conditions to select stocks.",
|
||||
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)
|
||||
// InvokableRun matches the current Python invocation result: valid arguments
|
||||
// produce an empty report and no error because the upstream call is disabled.
|
||||
func (w *WencaiTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
err := ctx.Err()
|
||||
return wencaiErrJSON(err), err
|
||||
default:
|
||||
}
|
||||
if p.Query == "" {
|
||||
return wencaiErrJSON(errors.New(wencaiUnsupportedMessage)),
|
||||
errors.New(wencaiUnsupportedMessage)
|
||||
|
||||
params := w.defaults
|
||||
if strings.TrimSpace(argsJSON) != "" {
|
||||
var runtimeParams wencaiParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &runtimeParams); err != nil {
|
||||
err = fmt.Errorf("wencai: parse arguments: %w", err)
|
||||
return wencaiErrJSON(err), err
|
||||
}
|
||||
params = mergeWencaiParams(params, runtimeParams)
|
||||
}
|
||||
return wencaiErrJSON(errors.New(wencaiUnsupportedMessage)),
|
||||
errors.New(wencaiUnsupportedMessage)
|
||||
_ = params
|
||||
return wencaiJSON(wencaiEnvelope{Report: ""}), nil
|
||||
}
|
||||
|
||||
func mergeWencaiParams(defaults, params wencaiParams) wencaiParams {
|
||||
if params.Query == "" {
|
||||
params.Query = defaults.Query
|
||||
}
|
||||
if params.TopN == 0 {
|
||||
params.TopN = defaults.TopN
|
||||
}
|
||||
if strings.TrimSpace(params.QueryType) == "" {
|
||||
params.QueryType = defaults.QueryType
|
||||
}
|
||||
return params
|
||||
}
|
||||
|
||||
func isWencaiQueryTypeSupported(queryType string) bool {
|
||||
_, ok := wencaiQueryTypes[queryType]
|
||||
return ok
|
||||
}
|
||||
|
||||
func wencaiJSON(env wencaiEnvelope) string {
|
||||
|
||||
@@ -19,54 +19,39 @@ package tool
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestWencai_StubsUnsupported(t *testing.T) {
|
||||
func TestWencai_InvokeMatchesCurrentPythonResult(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: `{}`,
|
||||
},
|
||||
{name: "query", args: `{"query":"商业航天"}`},
|
||||
{name: "empty query", args: `{"query":""}`},
|
||||
{name: "missing query", args: `{}`},
|
||||
{name: "empty arguments", 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())
|
||||
out, err := NewWencaiTool().InvokableRun(context.Background(), tc.args)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun errored: %v (out=%s)", err, out)
|
||||
}
|
||||
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 err := json.Unmarshal([]byte(out), &env); err != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", err, out)
|
||||
}
|
||||
if env.Error == "" {
|
||||
t.Errorf("env.Error = empty, want populated")
|
||||
if env.Report != "" {
|
||||
t.Fatalf("report = %q, want empty string", env.Report)
|
||||
}
|
||||
if !strings.Contains(env.Error, "同花顺") {
|
||||
t.Errorf("env.Error = %q, want to mention 同花顺", env.Error)
|
||||
if env.Error != "" {
|
||||
t.Fatalf("_ERROR = %q, want empty", env.Error)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -75,31 +60,137 @@ func TestWencai_StubsUnsupported(t *testing.T) {
|
||||
func TestWencai_RejectsMalformedJSON(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewWencaiTool()
|
||||
_, err := tool.InvokableRun(context.Background(), `{not json`)
|
||||
out, err := NewWencaiTool().InvokableRun(context.Background(), `{not json`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for malformed JSON, got nil")
|
||||
t.Fatal("expected malformed JSON error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "同花顺") {
|
||||
t.Errorf("err = %q, want to mention 同花顺", err.Error())
|
||||
if !strings.Contains(err.Error(), "parse arguments") {
|
||||
t.Fatalf("err = %q, want parse arguments", err.Error())
|
||||
}
|
||||
var env wencaiEnvelope
|
||||
if jsonErr := json.Unmarshal([]byte(out), &env); jsonErr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jsonErr, out)
|
||||
}
|
||||
if env.Error == "" {
|
||||
t.Fatal("_ERROR is empty, want parse error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_Info(t *testing.T) {
|
||||
func TestWencai_RespectsCanceledContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := NewWencaiTool()
|
||||
info, err := tool.Info(context.Background())
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
out, err := NewWencaiTool().InvokableRun(ctx, `{"query":"商业航天"}`)
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("err = %v, want context.Canceled", err)
|
||||
}
|
||||
var env wencaiEnvelope
|
||||
if jsonErr := json.Unmarshal([]byte(out), &env); jsonErr != nil {
|
||||
t.Fatalf("output is not valid JSON: %v (raw=%s)", jsonErr, out)
|
||||
}
|
||||
if env.Error != context.Canceled.Error() {
|
||||
t.Fatalf("_ERROR = %q, want %q", env.Error, context.Canceled.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_InfoMatchesPythonMeta(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
info, err := NewWencaiTool().Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != "wencai" {
|
||||
t.Errorf("Name = %q, want wencai", info.Name)
|
||||
if info.Name != "iwencai" {
|
||||
t.Fatalf("Name = %q, want iwencai", 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, "unsupported") {
|
||||
t.Fatalf("Desc still exposes stub behavior: %q", info.Desc)
|
||||
}
|
||||
if !strings.Contains(info.Desc, "STUB") && !strings.Contains(info.Desc, "Python") {
|
||||
t.Errorf("Desc = %q, want to flag stub status", info.Desc)
|
||||
schema, err := info.ParamsOneOf.ToJSONSchema()
|
||||
if err != nil {
|
||||
t.Fatalf("ToJSONSchema: %v", err)
|
||||
}
|
||||
raw, err := json.Marshal(schema)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal schema: %v", err)
|
||||
}
|
||||
params := string(raw)
|
||||
if !strings.Contains(params, `"query"`) || !strings.Contains(params, `"required":["query"]`) {
|
||||
t.Fatalf("schema does not require query: %s", params)
|
||||
}
|
||||
if strings.Contains(params, `"top_n"`) || strings.Contains(params, `"query_type"`) {
|
||||
t.Fatalf("schema leaked node parameters: %s", params)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_BuildByNameAcceptsNodeParams(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
built, err := BuildByName("wencai", map[string]any{
|
||||
"top_n": float64(20),
|
||||
"query_type": "fund",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildByName: %v", err)
|
||||
}
|
||||
wencai, ok := built.(*WencaiTool)
|
||||
if !ok {
|
||||
t.Fatalf("built type = %T, want *WencaiTool", built)
|
||||
}
|
||||
if wencai.defaults.TopN != 20 {
|
||||
t.Fatalf("defaults.TopN = %d, want 20", wencai.defaults.TopN)
|
||||
}
|
||||
if wencai.defaults.QueryType != "fund" {
|
||||
t.Fatalf("defaults.QueryType = %q, want fund", wencai.defaults.QueryType)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_BuildByNameUsesPythonDefaults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
built, err := BuildByName("wencai", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildByName: %v", err)
|
||||
}
|
||||
wencai := built.(*WencaiTool)
|
||||
if wencai.defaults.TopN != 10 || wencai.defaults.QueryType != "stock" {
|
||||
t.Fatalf("defaults = %+v, want top_n=10 query_type=stock", wencai.defaults)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_BuildByNameRejectsInvalidNodeParams(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
params map[string]any
|
||||
}{
|
||||
{name: "zero top_n", params: map[string]any{"top_n": 0}},
|
||||
{name: "fractional top_n", params: map[string]any{"top_n": 1.5}},
|
||||
{name: "string top_n", params: map[string]any{"top_n": "10"}},
|
||||
{name: "invalid query_type", params: map[string]any{"query_type": "crypto"}},
|
||||
{name: "non-string query_type", params: map[string]any{"query_type": 1}},
|
||||
{name: "unknown key", params: map[string]any{"page": 1}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if _, err := BuildByName("wencai", tc.params); err == nil {
|
||||
t.Fatalf("BuildByName(%v) succeeded, want validation error", tc.params)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWencai_MergeDefaults(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got := mergeWencaiParams(
|
||||
wencaiParams{Query: "configured", TopN: 20, QueryType: "fund"},
|
||||
wencaiParams{Query: "商业航天"},
|
||||
)
|
||||
if got.Query != "商业航天" || got.TopN != 20 || got.QueryType != "fund" {
|
||||
t.Fatalf("merged params = %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user