From 9289a183a21db769eb32f028d15fe43c017e7288 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Tue, 14 Jul 2026 14:57:22 +0800 Subject: [PATCH] fix: support keenable component (#16884) ### Summary As title image --- internal/agent/canvas/runner.go | 6 +- internal/agent/component/fixture_stubs.go | 1 + .../component/keenable_component_test.go | 169 +++++++++++++++ internal/agent/component/message.go | 2 +- internal/agent/component/message_test.go | 14 ++ .../agent/component/universe_a_wrappers.go | 197 ++++++++++++++++++ internal/agent/runtime/state.go | 71 +++++++ internal/agent/runtime/state_test.go | 23 ++ internal/agent/tool/keenable.go | 13 ++ internal/agent/tool/keenable_test.go | 10 + internal/service/agent.go | 43 ++-- 11 files changed, 532 insertions(+), 17 deletions(-) create mode 100644 internal/agent/component/keenable_component_test.go diff --git a/internal/agent/canvas/runner.go b/internal/agent/canvas/runner.go index 6380488728..138fd7a997 100644 --- a/internal/agent/canvas/runner.go +++ b/internal/agent/canvas/runner.go @@ -107,15 +107,15 @@ type NodeFinishedData struct { // MessageEvent is the JSON payload for Type=="message" frames. type MessageEvent struct { - Content string `json:"content"` - Reference []interface{} `json:"reference,omitempty"` + Content string `json:"content"` + Reference interface{} `json:"reference,omitempty"` } // MessageEndEvent is the JSON payload for Type=="message_end" frames. type MessageEndEvent struct { Status *string `json:"status,omitempty"` Attachment []interface{} `json:"attachment,omitempty"` - Reference []interface{} `json:"reference,omitempty"` + Reference interface{} `json:"reference,omitempty"` } // WaitingForUserEvent is the JSON payload for Type=="waiting_for_user" diff --git a/internal/agent/component/fixture_stubs.go b/internal/agent/component/fixture_stubs.go index 1aa94f12f5..c60d49ace5 100644 --- a/internal/agent/component/fixture_stubs.go +++ b/internal/agent/component/fixture_stubs.go @@ -521,6 +521,7 @@ func init() { Register("GitHub", newGitHubComponent) Register("Wikipedia", newWikipediaComponent) Register("DuckDuckGo", newDuckDuckGoComponent) + Register("KeenableSearch", newKeenableSearchComponent) Register("Google", newGoogleComponent) Register("GoogleScholar", newGoogleScholarComponent) Register("ArXiv", newArxivComponent) diff --git a/internal/agent/component/keenable_component_test.go b/internal/agent/component/keenable_component_test.go new file mode 100644 index 0000000000..03886bb28f --- /dev/null +++ b/internal/agent/component/keenable_component_test.go @@ -0,0 +1,169 @@ +// +// 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" + "strings" + "testing" + + einotool "github.com/cloudwego/eino/components/tool" + + "ragflow/internal/agent/runtime" +) + +type fakeKeenableInvoker struct { + args map[string]any +} + +func (f *fakeKeenableInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) { + if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil { + return "", err + } + return `{"results":[{"title":"Keenable result","url":"https://example.com/item","description":"Fresh search result"}]}`, nil +} + +func TestKeenableSearch_RegisteredFactory(t *testing.T) { + t.Parallel() + + c, err := New("KeenableSearch", map[string]any{"api_key": "key-1", "mode": "realtime", "top_n": float64(3)}) + if err != nil { + t.Fatalf("New(KeenableSearch) errored: %v", err) + } + kc, ok := c.(*keenableSearchComponent) + if !ok { + t.Fatalf("New(KeenableSearch) returned %T, want *keenableSearchComponent", c) + } + if kc.apiKey != "key-1" { + t.Fatalf("apiKey = %q, want key-1", kc.apiKey) + } + if kc.mode != "realtime" { + t.Fatalf("mode = %q, want realtime", kc.mode) + } + if kc.topN != 3 { + t.Fatalf("topN = %d, want 3", kc.topN) + } + + formGetter, ok := c.(interface{ GetInputForm() map[string]any }) + if !ok { + t.Fatal("KeenableSearch component does not expose GetInputForm") + } + form := formGetter.GetInputForm() + query, ok := form["query"].(map[string]any) + if !ok { + t.Fatalf("GetInputForm()[query] has type %T, want map", form["query"]) + } + if query["type"] != "line" { + t.Fatalf("GetInputForm()[query][type] = %v, want line", query["type"]) + } + site, ok := form["site"].(map[string]any) + if !ok { + t.Fatalf("GetInputForm()[site] has type %T, want map", form["site"]) + } + if site["type"] != "line" { + t.Fatalf("GetInputForm()[site][type] = %v, want line", site["type"]) + } + if _, ok := c.Outputs()["formalized_content"]; !ok { + t.Fatal("Outputs() missing formalized_content") + } + if _, ok := c.Outputs()["json"]; !ok { + t.Fatal("Outputs() missing json") + } +} + +func TestKeenableSearch_InvokeAdaptsCanvasInputsAndOutputs(t *testing.T) { + t.Parallel() + + fake := &fakeKeenableInvoker{} + c := newKeenableSearchComponentWithInvoker(fake, map[string]any{ + "mode": "pro", + "top_n": 5, + "site": "example.com", + }) + + state := runtime.NewCanvasState("run-keenable", "task-keenable") + ctx := runtime.WithState(context.Background(), state) + + out, err := c.Invoke(ctx, map[string]any{ + "query": " agent search ", + "mode": "realtime", + "top_n": float64(2), + }) + if err != nil { + t.Fatalf("Invoke errored: %v", err) + } + + if got := fake.args["query"]; got != "agent search" { + t.Errorf("query arg = %v, want trimmed query", got) + } + if got := fake.args["mode"]; got != "realtime" { + t.Errorf("mode arg = %v, want realtime", got) + } + if got := fake.args["top_n"]; got != float64(2) && got != 2 { + t.Errorf("top_n arg = %v, want 2", got) + } + if got := fake.args["site"]; got != "example.com" { + t.Errorf("site arg = %v, want default site", got) + } + + formalized, _ := out["formalized_content"].(string) + for _, want := range []string{"Keenable result", "https://example.com/item", "Fresh search result"} { + if !strings.Contains(formalized, want) { + t.Errorf("formalized_content missing %q: %s", want, formalized) + } + } + + results, ok := out["json"].([]any) + if !ok { + t.Fatalf("json output has type %T, want []any", out["json"]) + } + if len(results) != 1 { + t.Fatalf("json output length = %d, want 1", len(results)) + } + + reference := state.GetRetrievalReference() + chunks, _ := reference["chunks"].([]any) + if len(chunks) != 1 { + t.Fatalf("reference chunks length = %d, want 1", len(chunks)) + } + chunk, _ := chunks[0].(map[string]any) + if chunk["document_name"] != "Keenable result" || chunk["url"] != "https://example.com/item" { + t.Fatalf("reference chunk metadata = %#v", chunk) + } + docAggs, _ := reference["doc_aggs"].([]any) + if len(docAggs) != 1 { + t.Fatalf("reference doc_aggs length = %d, want 1", len(docAggs)) + } +} + +func TestKeenableSearch_InvokeEmptyQueryReturnsEmptyPayload(t *testing.T) { + t.Parallel() + + c := newKeenableSearchComponentWithInvoker(&fakeKeenableInvoker{}, nil) + out, err := c.Invoke(context.Background(), map[string]any{"query": " "}) + if err != nil { + t.Fatalf("Invoke errored: %v", err) + } + if got := out["formalized_content"]; got != "" { + t.Errorf("formalized_content = %v, want empty string", got) + } + results, ok := out["json"].([]any) + if !ok || len(results) != 0 { + t.Fatalf("json output = %#v, want empty []any", out["json"]) + } +} diff --git a/internal/agent/component/message.go b/internal/agent/component/message.go index 426f79b9f0..5a834b8b84 100644 --- a/internal/agent/component/message.go +++ b/internal/agent/component/message.go @@ -175,7 +175,7 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m return nil, fmt.Errorf("Message: nil canvas state") } - text, _ := inputs["text"].(string) + text := extractMessageText(inputs) if text == "" { text = m.text } diff --git a/internal/agent/component/message_test.go b/internal/agent/component/message_test.go index bdede452a9..5983bb5fa1 100644 --- a/internal/agent/component/message_test.go +++ b/internal/agent/component/message_test.go @@ -98,6 +98,20 @@ func TestMessage_NoTemplate(t *testing.T) { } } +func TestMessage_RuntimeContentInput(t *testing.T) { + c, _ := NewMessageComponent(nil) + state := canvas.NewCanvasState("run-4", "task-4") + ctx := withStateForTest(context.Background(), state) + + out, err := c.Invoke(ctx, map[string]any{"content": "from upstream", "stream": false}) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if got, _ := out["content"].(string); got != "from upstream" { + t.Errorf("content: got %q, want %q", got, "from upstream") + } +} + // withStateForTest is a thin alias for canvas.WithState kept for // readability at the test call sites (also defined in begin_test.go). // Same package → same symbol; defining it twice is a compile error, diff --git a/internal/agent/component/universe_a_wrappers.go b/internal/agent/component/universe_a_wrappers.go index cf98059dcb..fa1123815e 100644 --- a/internal/agent/component/universe_a_wrappers.go +++ b/internal/agent/component/universe_a_wrappers.go @@ -317,6 +317,10 @@ type duckDuckGoInvoker interface { InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) } +type keenableInvoker interface { + InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) +} + // bgptComponent delegates to internal/agent/tool/BGPTTool and adapts // the tool envelope to the BGPT canvas output contract. type bgptComponent struct { @@ -574,6 +578,199 @@ func (c *duckDuckGoComponent) Stream(_ context.Context, _ map[string]any) (<-cha return nil, nil } +// keenableSearchComponent delegates to internal/agent/tool/KeenableTool. +type keenableSearchComponent struct { + inner keenableInvoker + apiKey string + mode string + topN int + site string +} + +func newKeenableSearchComponent(params map[string]any) (Component, error) { + apiKey := stringParam(params["api_key"]) + return newKeenableSearchComponentWithInvoker(agenttool.NewKeenableToolWithAPIKey(nil, apiKey), params), nil +} + +func newKeenableSearchComponentWithInvoker(inner keenableInvoker, params map[string]any) Component { + if inner == nil { + inner = agenttool.NewKeenableTool() + } + mode := strings.TrimSpace(stringParam(params["mode"])) + if mode == "" { + mode = "pro" + } + topN := toIntParam(params["top_n"]) + if topN <= 0 { + topN = 10 + } + return &keenableSearchComponent{ + inner: inner, + apiKey: stringParam(params["api_key"]), + mode: mode, + topN: topN, + site: stringParam(params["site"]), + } +} + +func (c *keenableSearchComponent) Name() string { return "KeenableSearch" } + +func (c *keenableSearchComponent) Inputs() map[string]string { + return map[string]string{ + "query": "Search query.", + "site": "Optional single-domain filter.", + "api_key": "Optional Keenable API key.", + "mode": "Search mode: pro or realtime.", + "top_n": "Maximum number of results.", + } +} + +func (c *keenableSearchComponent) GetInputForm() map[string]any { + return map[string]any{ + "query": map[string]any{ + "name": "Query", + "type": "line", + }, + "site": map[string]any{ + "name": "Site", + "type": "line", + }, + } +} + +func (c *keenableSearchComponent) Outputs() map[string]string { + return map[string]string{ + "formalized_content": "Rendered search results for downstream LLM prompts.", + "json": "Raw Keenable result list.", + } +} + +func (c *keenableSearchComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { + query := strings.TrimSpace(stringParam(inputs["query"])) + if query == "" { + return map[string]any{"formalized_content": "", "json": []any{}}, nil + } + + args := map[string]any{ + "query": query, + "mode": c.mode, + "top_n": c.topN, + } + if mode := strings.TrimSpace(stringParam(inputs["mode"])); mode != "" { + args["mode"] = mode + } + if topN := toIntParam(inputs["top_n"]); topN > 0 { + args["top_n"] = topN + } + site := strings.TrimSpace(stringParam(inputs["site"])) + if site == "" { + site = strings.TrimSpace(c.site) + } + if site != "" { + args["site"] = site + } + + invoker := c.inner + if apiKey := strings.TrimSpace(stringParam(inputs["api_key"])); apiKey != "" && apiKey != strings.TrimSpace(c.apiKey) { + invoker = agenttool.NewKeenableToolWithAPIKey(nil, apiKey) + } + + argsJSON, _ := json.Marshal(args) + out, err := invoker.InvokableRun(ctx, string(argsJSON)) + decoded := parseToolEnvelope(out) + results := anySlice(decoded["results"]) + if existing, _ := decoded["_ERROR"].(string); strings.TrimSpace(existing) != "" { + return map[string]any{ + "formalized_content": "", + "json": results, + "_ERROR": existing, + }, nil + } + if err != nil { + if len(decoded) > 0 { + return map[string]any{ + "formalized_content": "", + "json": results, + "_ERROR": decoded["_ERROR"], + }, nil + } + return nil, fmt.Errorf("canvas: KeenableSearch: %w", err) + } + chunks, docAggs := buildKeenableReferences(results) + if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil { + state.SetRetrievalReferences(chunks, docAggs) + } + return map[string]any{ + "formalized_content": renderKeenableReferences(chunks), + "json": results, + }, nil +} + +func (c *keenableSearchComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { + return nil, nil +} + +func buildKeenableReferences(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 _, item := range results { + m, ok := item.(map[string]any) + if !ok { + continue + } + content := truncateRunes(strings.TrimSpace(keenableValueString(m["description"])), 10000) + if content == "" || content == "None" { + continue + } + documentID := strconv.FormatInt(githubHashInt(content, 100000000), 10) + title := keenableValueString(m["title"]) + url := keenableValueString(m["url"]) + displayID := strconv.FormatInt(githubHashInt(documentID, 500), 10) + chunks = append(chunks, map[string]any{ + "id": displayID, + "chunk_id": documentID, + "content": content, + "doc_id": documentID, + "document_id": documentID, + "docnm_kwd": title, + "document_name": title, + "similarity": 1, + "score": 1, + "url": url, + }) + docAggs = append(docAggs, map[string]any{ + "doc_name": title, + "doc_id": documentID, + "count": 1, + "url": url, + }) + } + return chunks, docAggs +} + +func renderKeenableReferences(chunks []map[string]any) string { + if len(chunks) == 0 { + return "" + } + blocks := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + blocks = append(blocks, strings.Join([]string{ + "\nID: " + keenableValueString(chunk["id"]), + "├── Title: " + keenableValueString(chunk["docnm_kwd"]), + "├── URL: " + keenableValueString(chunk["url"]), + "└── Content:\n" + keenableValueString(chunk["content"]), + }, "\n")) + } + return strings.Join(blocks, "\n") +} + +func keenableValueString(value any) string { + if value == nil { + return "None" + } + return strings.TrimSpace(fmt.Sprint(value)) +} + func renderDuckDuckGoResults(results []any) string { if len(results) == 0 { return "" diff --git a/internal/agent/runtime/state.go b/internal/agent/runtime/state.go index 7413f4f241..7d1ee9d977 100644 --- a/internal/agent/runtime/state.go +++ b/internal/agent/runtime/state.go @@ -33,6 +33,7 @@ package runtime import ( "encoding/json" "fmt" + "sort" "strings" "sync" "sync/atomic" @@ -368,6 +369,76 @@ func (s *CanvasState) GetRetrievalChunks() []map[string]any { return out } +// GetRetrievalReference returns the run-level reference payload consumed by +// the agent chat stream. It mirrors Python canvas.py's message_end.reference +// shape while keeping doc_aggs as a list for the current Go frontend path. +func (s *CanvasState) GetRetrievalReference() map[string]any { + if s == nil { + return nil + } + s.mu.RLock() + defer s.mu.RUnlock() + if len(s.Retrieval) == 0 { + return nil + } + + chunks := copyRetrievalList(s.Retrieval["chunks"]) + docAggs := copyRetrievalDocAggs(s.Retrieval["doc_aggs"]) + if len(chunks) == 0 && len(docAggs) == 0 { + return nil + } + return map[string]any{ + "chunks": chunks, + "doc_aggs": docAggs, + "total": len(chunks), + } +} + +func copyRetrievalList(value any) []any { + switch list := value.(type) { + case []any: + out := make([]any, len(list)) + copy(out, list) + return out + case []map[string]any: + out := make([]any, 0, len(list)) + for _, item := range list { + out = append(out, item) + } + return out + default: + return nil + } +} + +func copyRetrievalDocAggs(value any) []any { + switch aggs := value.(type) { + case []any: + out := make([]any, len(aggs)) + copy(out, aggs) + return out + case []map[string]any: + out := make([]any, 0, len(aggs)) + for _, item := range aggs { + out = append(out, item) + } + return out + case map[string]any: + keys := make([]string, 0, len(aggs)) + for key := range aggs { + keys = append(keys, key) + } + sort.Strings(keys) + out := make([]any, 0, len(keys)) + for _, key := range keys { + out = append(out, aggs[key]) + } + return out + default: + return nil + } +} + // SetRetrievalChunks records the supplied chunks into // state.Retrieval["chunks"]. Existing entries are replaced // (last-writer-wins) so a multi-tool canvas reflects the most diff --git a/internal/agent/runtime/state_test.go b/internal/agent/runtime/state_test.go index 2078074650..ed8e8f3b0d 100644 --- a/internal/agent/runtime/state_test.go +++ b/internal/agent/runtime/state_test.go @@ -133,6 +133,29 @@ func TestCanvasState_SetRetrievalReferencesMergesCalls(t *testing.T) { } } +func TestCanvasState_GetRetrievalReferenceReturnsFrontendPayload(t *testing.T) { + t.Parallel() + + state := NewCanvasState("run-1", "task-1") + state.SetRetrievalReferences( + []map[string]any{{"id": "chunk-1", "document_name": "Doc 1"}}, + []map[string]any{{"doc_name": "Doc 1", "doc_id": "doc-1", "count": 1}}, + ) + + reference := state.GetRetrievalReference() + chunks, _ := reference["chunks"].([]any) + if len(chunks) != 1 { + t.Fatalf("chunks length = %d, want 1", len(chunks)) + } + docAggs, _ := reference["doc_aggs"].([]any) + if len(docAggs) != 1 { + t.Fatalf("doc_aggs length = %d, want 1", len(docAggs)) + } + if reference["total"] != 1 { + t.Fatalf("total = %v, want 1", reference["total"]) + } +} + func contains(s, substr string) bool { for i := 0; i+len(substr) <= len(s); i++ { if s[i:i+len(substr)] == substr { diff --git a/internal/agent/tool/keenable.go b/internal/agent/tool/keenable.go index 5d946f6b69..aeecfe8f2b 100644 --- a/internal/agent/tool/keenable.go +++ b/internal/agent/tool/keenable.go @@ -209,6 +209,19 @@ func (k *KeenableTool) Info(_ context.Context) (*schema.ToolInfo, error) { }, nil } +func (k *KeenableTool) GetInputForm() map[string]any { + return map[string]any{ + "query": map[string]any{ + "name": "Query", + "type": "line", + }, + "site": map[string]any{ + "name": "Site", + "type": "line", + }, + } +} + // InvokableRun performs the Keenable search. func (k *KeenableTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) { var p keenableParams diff --git a/internal/agent/tool/keenable_test.go b/internal/agent/tool/keenable_test.go index c0fd3bcaf8..82a12213d3 100644 --- a/internal/agent/tool/keenable_test.go +++ b/internal/agent/tool/keenable_test.go @@ -389,4 +389,14 @@ func TestKeenable_Info(t *testing.T) { if strings.Contains(string(paramsJSON), "api_key") { t.Fatalf("Info ParamsOneOf unexpectedly exposes api_key: %s", string(paramsJSON)) } + form := tool.GetInputForm() + for _, key := range []string{"query", "site"} { + field, ok := form[key].(map[string]any) + if !ok { + t.Fatalf("GetInputForm()[%s] = %#v, want field map", key, form[key]) + } + if field["type"] != "line" { + t.Fatalf("GetInputForm()[%s][type] = %v, want line", key, field["type"]) + } + } } diff --git a/internal/service/agent.go b/internal/service/agent.go index 99a2e8ab89..03174f4be6 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -1115,7 +1115,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv // node_finished events are already emitted per-node by the // statePost wrappers in scheduler.go. var answer string - var reference []interface{} + var legacyReference []interface{} var downloads any now := float64(time.Now().UnixNano()) / 1e9 for _, bucket := range state.Snapshot() { @@ -1131,12 +1131,13 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv answer = v } if v, ok := bucket["reference"].([]interface{}); ok { - reference = append(reference, v...) + legacyReference = append(legacyReference, v...) } if v, ok := bucket["downloads"]; ok && !emptyDownloadValue(v) { downloads = v } } + referencePayload := agentRunReferencePayload(state, legacyReference) if err != nil { common.Debug("RunAgent invoke err", @@ -1149,30 +1150,30 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv if canvas.IsInterruptError(err) { s.markRunFailed(ctx2, runID, "interrupt: "+err.Error()) if answer != "" { - s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference) + s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload) msgData, _ := json.Marshal(canvas.MessageEvent{ Content: answer, - Reference: reference, + Reference: referencePayload, }) emit("message", string(msgData)) meData, _ := json.Marshal(canvas.MessageEndEvent{ - Reference: reference, + Reference: referencePayload, }) emit("message_end", string(meData)) } return state, err } if shouldTreatAsCompletedLoopRun(err, answer) { - s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference) + s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload) msgData, _ := json.Marshal(canvas.MessageEvent{ Content: answer, - Reference: reference, + Reference: referencePayload, }) emit("message", string(msgData)) meData, _ := json.Marshal(canvas.MessageEndEvent{ - Reference: reference, + Reference: referencePayload, }) emit("message_end", string(meData)) @@ -1196,15 +1197,15 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv } // Emit message + message_end (mirrors Python's ans dict). - s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference) + s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload) msgData, _ := json.Marshal(canvas.MessageEvent{ Content: answer, - Reference: reference, + Reference: referencePayload, }) emit("message", string(msgData)) meData, _ := json.Marshal(canvas.MessageEndEvent{ - Reference: reference, + Reference: referencePayload, }) emit("message_end", string(meData)) @@ -1261,7 +1262,7 @@ func emptyDownloadValue(value any) bool { } } -func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID string, userInput any, answer string, reference []interface{}) { +func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID string, userInput any, answer string, reference map[string]interface{}) { if sessionID == "" || s == nil || s.api4ConversationDAO == nil || dao.DB == nil { return } @@ -1283,7 +1284,7 @@ func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID stri session.Message = raw } references := parseAgentSessionReferences(session.Reference) - references = append(references, normalizeAgentReferenceEntry(map[string]interface{}{"chunks": reference})) + references = append(references, normalizeAgentReferenceEntry(reference)) if raw, err := json.Marshal(references); err == nil { session.Reference = raw } @@ -1292,6 +1293,22 @@ func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID stri } } +func agentRunReferencePayload(state *canvas.CanvasState, legacyChunks []interface{}) map[string]interface{} { + if state != nil { + if reference := state.GetRetrievalReference(); len(reference) > 0 { + return reference + } + } + if len(legacyChunks) == 0 { + return nil + } + return map[string]interface{}{ + "chunks": legacyChunks, + "doc_aggs": []interface{}{}, + "total": len(legacyChunks), + } +} + func stringifyAgentUserInput(userInput any) string { switch v := userInput.(type) { case nil: