From 59ace30b60936ca09f6338516b39e412a29b7717 Mon Sep 17 00:00:00 2001 From: shihyunhuang Date: Sat, 6 Jun 2026 16:36:10 -0700 Subject: [PATCH 1/4] dev knowledge.go structure Signed-off-by: Pete Chen --- core/services/agents/executor.go | 8 ++- core/services/agents/knowledge.go | 100 +++++++++++++++++++++++------- 2 files changed, 84 insertions(+), 24 deletions(-) diff --git a/core/services/agents/executor.go b/core/services/agents/executor.go index 81481600c6a2..2f5821e87a70 100644 --- a/core/services/agents/executor.go +++ b/core/services/agents/executor.go @@ -168,10 +168,12 @@ func ExecuteChatWithLLM(ctx context.Context, llm cogito.LLM, cfg *AgentConfig, m } if cfg.EnableKnowledgeBase && (kbMode == KBModeAutoSearch || kbMode == KBModeBoth) { - kbResults := KBAutoSearchPrompt(ctx, effectiveURL, effectiveKey, cfg.Name, message, cfg.KnowledgeBaseResults, userID) - if kbResults != "" { - fragment = fragment.AddMessage(cogito.SystemMessageRole, kbResults) + kbResult := KBAutoSearchPrompt(ctx, effectiveURL, effectiveKey, cfg.Name, message, cfg.KnowledgeBaseResults, userID) + if kbResult.Prompt != "" { + fragment = fragment.AddMessage(cogito.SystemMessageRole, kbResult.Prompt) } + // NOTE: kbResult.Citations is the structured source list. Appending it as + // a Sources block to the answer is person 2's work in this file. } // User message diff --git a/core/services/agents/knowledge.go b/core/services/agents/knowledge.go index 8c0afb012984..54ae57be5167 100644 --- a/core/services/agents/knowledge.go +++ b/core/services/agents/knowledge.go @@ -17,10 +17,19 @@ import ( "github.com/mudler/LocalAI/pkg/httpclient" ) +// Metadata keys populated by localrecall for every stored chunk. The original +// upload file name lives under file_name (used for display); source holds the +// collection entry key ("/") used to build the raw-file URL. +const ( + kbMetadataFileName = "file_name" + kbMetadataSource = "source" +) + // KBSearchResult represents a search result from the knowledge base. +// Field names mirror the collection search endpoint's JSON response. type KBSearchResult struct { Content string `json:"content"` - Score float64 `json:"score"` + ID string `json:"id"` Similarity float64 `json:"similarity"` Metadata map[string]string `json:"metadata"` } @@ -31,12 +40,32 @@ type kbSearchResponse struct { Count int `json:"count"` } -// KBAutoSearchPrompt queries the knowledge base with the user's message -// and returns a system prompt block with relevant results. +// KBCitation is a single source document that a KB search drew from. Citations +// travel alongside the prompt as structured data so the consumer (and UI) can +// render clickable source links, independent of what the model writes inline. +type KBCitation struct { + // FileName is the original uploaded file name, for display (e.g. "report.pdf"). + FileName string `json:"file_name"` + // EntryKey is the collection entry identifier ("/"), used to + // build the raw-file URL and as the de-duplication key. + EntryKey string `json:"entry_key"` +} + +// KBSearchContext is the result of an auto-search against the knowledge base: +// the system-prompt block to feed the model, plus the de-duplicated list of +// source documents the results were drawn from. +type KBSearchContext struct { + Prompt string `json:"prompt"` + Citations []KBCitation `json:"citations"` +} + +// KBAutoSearchPrompt queries the knowledge base with the user's message and +// returns a KBSearchContext: a system prompt block with the relevant results +// plus the de-duplicated source citations those results came from. // Uses LocalAI's collection search endpoint via the API. -func KBAutoSearchPrompt(ctx context.Context, apiURL, apiKey, collection, query string, maxResults int, userID string) string { +func KBAutoSearchPrompt(ctx context.Context, apiURL, apiKey, collection, query string, maxResults int, userID string) KBSearchContext { if collection == "" || query == "" { - return "" + return KBSearchContext{} } if maxResults <= 0 { @@ -56,7 +85,7 @@ func KBAutoSearchPrompt(ctx context.Context, apiURL, apiKey, collection, query s req, err := http.NewRequestWithContext(ctx, http.MethodPost, searchURL, strings.NewReader(string(reqBody))) if err != nil { xlog.Warn("KB auto-search: failed to create request", "error", err) - return "" + return KBSearchContext{} } req.Header.Set("Content-Type", "application/json") if apiKey != "" { @@ -66,41 +95,70 @@ func KBAutoSearchPrompt(ctx context.Context, apiURL, apiKey, collection, query s resp, err := httpclient.New().Do(req) if err != nil { xlog.Warn("KB auto-search: request failed", "error", err) - return "" + return KBSearchContext{} } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) xlog.Warn("KB auto-search: non-200 response", "status", resp.StatusCode, "body", string(body)) - return "" + return KBSearchContext{} } var searchResp kbSearchResponse if err := json.NewDecoder(resp.Body).Decode(&searchResp); err != nil { xlog.Warn("KB auto-search: failed to decode response", "error", err) - return "" + return KBSearchContext{} } if len(searchResp.Results) == 0 { - return "" + return KBSearchContext{} } - // Format results as a system prompt block (same format as LocalAGI) + // Build the system prompt block, labelling each chunk with its source file + // so the model can attribute inline, and collect the structured citations. var sb strings.Builder sb.WriteString("Given the user input you have the following in memory:\n") - for i, r := range searchResp.Results { - sb.WriteString(fmt.Sprintf("- %s", r.Content)) - if len(r.Metadata) > 0 { - meta, _ := json.Marshal(r.Metadata) - sb.WriteString(fmt.Sprintf(" (%s)", string(meta))) + + var citations []KBCitation + seen := make(map[string]struct{}) + + for _, r := range searchResp.Results { + fileName := r.Metadata[kbMetadataFileName] + source := r.Metadata[kbMetadataSource] + + label := fileName + if label == "" { + label = "unknown" } - if i < len(searchResp.Results)-1 { - sb.WriteString("\n") + sb.WriteString(fmt.Sprintf("[Source: %s]\n%s\n", label, r.Content)) + + // Citations are de-duplicated per source document: many chunks from the + // same file share one source key, so a file is listed only once. Skip + // results with no source key — they cannot be linked back to a document. + dedupKey := source + if dedupKey == "" { + dedupKey = fileName } + if dedupKey == "" { + continue + } + if _, ok := seen[dedupKey]; ok { + continue + } + seen[dedupKey] = struct{}{} + citations = append(citations, KBCitation{ + FileName: fileName, + EntryKey: source, + }) } - return sb.String() + sb.WriteString("When answering, cite sources using [Source: filename].") + + return KBSearchContext{ + Prompt: sb.String(), + Citations: citations, + } } // KBSearchMemoryArgs defines the arguments for the search_memory tool. @@ -121,10 +179,10 @@ func (t KBSearchMemoryTool) Run(args KBSearchMemoryArgs) (string, any, error) { ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() result := KBAutoSearchPrompt(ctx, t.APIURL, t.APIKey, t.Collection, args.Query, t.MaxResults, t.UserID) - if result == "" { + if result.Prompt == "" { return "No results found.", nil, nil } - return result, nil, nil + return result.Prompt, nil, nil } // KBAddMemoryArgs defines the arguments for the add_memory tool. From 55cab3cf9c934277efc3b1fed65fd577379f4a01 Mon Sep 17 00:00:00 2001 From: TLoE419 Date: Sat, 6 Jun 2026 17:50:12 -0700 Subject: [PATCH 2/4] feat(agents): append KB source citations to responses Render structured KB citations as a Sources block after agent responses, linking each source to the existing raw collection entry endpoint. Keep long-term memory writes on the original model response so citation blocks do not get stored back into the knowledge base. Tested with: go test ./core/services/agents Assisted-by: Codex:gpt-5 Signed-off-by: Pete Chen --- core/services/agents/executor.go | 105 ++++++++++++++++++++++- core/services/agents/executor_test.go | 117 ++++++++++++++++++++++++++ 2 files changed, 219 insertions(+), 3 deletions(-) diff --git a/core/services/agents/executor.go b/core/services/agents/executor.go index 2f5821e87a70..3d23da0c81bf 100644 --- a/core/services/agents/executor.go +++ b/core/services/agents/executor.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "fmt" + "net/url" "strings" "time" @@ -167,13 +168,13 @@ func ExecuteChatWithLLM(ctx context.Context, llm cogito.LLM, cfg *AgentConfig, m } } + var kbCitations []KBCitation if cfg.EnableKnowledgeBase && (kbMode == KBModeAutoSearch || kbMode == KBModeBoth) { kbResult := KBAutoSearchPrompt(ctx, effectiveURL, effectiveKey, cfg.Name, message, cfg.KnowledgeBaseResults, userID) if kbResult.Prompt != "" { fragment = fragment.AddMessage(cogito.SystemMessageRole, kbResult.Prompt) + kbCitations = kbResult.Citations } - // NOTE: kbResult.Citations is the structured source list. Appending it as - // a Sources block to the answer is person 2's work in this file. } // User message @@ -338,6 +339,8 @@ func ExecuteChatWithLLM(ctx context.Context, llm cogito.LLM, cfg *AgentConfig, m if cfg.StripThinkingTags && response != "" { response = stripThinkingTags(response) } + responseForMemory := response + response = appendKBCitations(response, cfg.Name, userID, kbCitations) // Save conversation to KB when long-term memory is enabled. // Use a detached context: the parent ctx may be cancelled (e.g. in distributed @@ -346,7 +349,7 @@ func ExecuteChatWithLLM(ctx context.Context, llm cogito.LLM, cfg *AgentConfig, m go func() { ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - saveConversationToKB(ctx, llm, effectiveURL, effectiveKey, cfg, message, response, userID) + saveConversationToKB(ctx, llm, effectiveURL, effectiveKey, cfg, message, responseForMemory, userID) }() } @@ -382,6 +385,102 @@ func isInternalCogitoTool(name string) bool { return false } +func appendKBCitations(response, collection, userID string, citations []KBCitation) string { + if strings.TrimSpace(response) == "" || len(citations) == 0 { + return response + } + + var lines []string + seen := make(map[string]struct{}) + for _, citation := range citations { + key := strings.TrimSpace(citation.EntryKey) + if key == "" { + key = strings.TrimSpace(citation.FileName) + } + if key == "" { + continue + } + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + + displayName := kbCitationDisplayName(citation) + if displayName == "" { + continue + } + + sourceURL := kbCitationRawFileURL(collection, citation.EntryKey, userID) + number := len(lines) + 1 + if sourceURL == "" { + lines = append(lines, fmt.Sprintf("[%d] %s", number, displayName)) + continue + } + lines = append(lines, fmt.Sprintf("[%d] [%s](%s)", number, escapeMarkdownLinkText(displayName), sourceURL)) + } + + if len(lines) == 0 { + return response + } + + var sb strings.Builder + sb.WriteString(strings.TrimRight(response, "\n")) + sb.WriteString("\n\nSources:\n") + for _, line := range lines { + sb.WriteString(line) + sb.WriteString("\n") + } + return strings.TrimRight(sb.String(), "\n") +} + +func kbCitationDisplayName(citation KBCitation) string { + if fileName := strings.TrimSpace(citation.FileName); fileName != "" { + return fileName + } + + segments := strings.Split(strings.Trim(strings.TrimSpace(citation.EntryKey), "/"), "/") + for i := len(segments) - 1; i >= 0; i-- { + if segment := strings.TrimSpace(segments[i]); segment != "" { + return segment + } + } + return "" +} + +func kbCitationRawFileURL(collection, entryKey, userID string) string { + collection = strings.TrimSpace(collection) + entryKey = strings.Trim(strings.TrimSpace(entryKey), "/") + if collection == "" || entryKey == "" { + return "" + } + + var escapedEntrySegments []string + for _, segment := range strings.Split(entryKey, "/") { + if segment == "" { + continue + } + escapedEntrySegments = append(escapedEntrySegments, url.PathEscape(segment)) + } + if len(escapedEntrySegments) == 0 { + return "" + } + + sourceURL := "/api/agents/collections/" + url.PathEscape(collection) + "/entries-raw/" + strings.Join(escapedEntrySegments, "/") + if userID != "" { + query := url.Values{} + query.Set("user_id", userID) + sourceURL += "?" + query.Encode() + } + return sourceURL +} + +func escapeMarkdownLinkText(text string) string { + text = strings.ReplaceAll(text, `\`, `\\`) + text = strings.ReplaceAll(text, "[", `\[`) + text = strings.ReplaceAll(text, "]", `\]`) + return text +} + // stripThinkingTags removes ... blocks from content. func stripThinkingTags(content string) string { // Simple implementation — remove content between and diff --git a/core/services/agents/executor_test.go b/core/services/agents/executor_test.go index 04d44462ca3c..5b86005fc3da 100644 --- a/core/services/agents/executor_test.go +++ b/core/services/agents/executor_test.go @@ -2,6 +2,8 @@ package agents import ( "context" + "net/http" + "net/http/httptest" "sync" "sync/atomic" @@ -73,6 +75,74 @@ var _ = DescribeTable("stripThinkingTags", Entry("adjacent tag pairs", "ab", ""), ) +var _ = DescribeTable("appendKBCitations", + func(response, collection, userID string, citations []KBCitation, want string) { + Expect(appendKBCitations(response, collection, userID, citations)).To(Equal(want)) + }, + Entry("leaves responses without citations unchanged", + "answer", + "agent", + "", + nil, + "answer", + ), + Entry("leaves blank responses unchanged", + "", + "agent", + "", + []KBCitation{{FileName: "source.pdf", EntryKey: "uuid/source.pdf"}}, + "", + ), + Entry("appends clickable source links", + "answer", + "my-agent", + "", + []KBCitation{{FileName: "new feature.pdf", EntryKey: "uuid/new feature.pdf"}}, + "answer\n\nSources:\n[1] [new feature.pdf](/api/agents/collections/my-agent/entries-raw/uuid/new%20feature.pdf)", + ), + Entry("deduplicates citations by entry key", + "answer", + "agent", + "", + []KBCitation{ + {FileName: "first.pdf", EntryKey: "uuid/shared.pdf"}, + {FileName: "second.pdf", EntryKey: "uuid/shared.pdf"}, + }, + "answer\n\nSources:\n[1] [first.pdf](/api/agents/collections/agent/entries-raw/uuid/shared.pdf)", + ), + Entry("uses plain text when entry key is missing", + "answer", + "agent", + "", + []KBCitation{{FileName: "source.pdf"}}, + "answer\n\nSources:\n[1] source.pdf", + ), + Entry("uses entry basename when filename is missing", + "answer", + "agent", + "", + []KBCitation{{EntryKey: "uuid/source.pdf"}}, + "answer\n\nSources:\n[1] [source.pdf](/api/agents/collections/agent/entries-raw/uuid/source.pdf)", + ), + Entry("adds user id query when present", + "answer", + "agent", + "user 1", + []KBCitation{{FileName: "source.pdf", EntryKey: "uuid/source.pdf"}}, + "answer\n\nSources:\n[1] [source.pdf](/api/agents/collections/agent/entries-raw/uuid/source.pdf?user_id=user+1)", + ), + Entry("escapes collection, path segments, and markdown link text", + "answer", + "agent one", + "", + []KBCitation{{FileName: "source [draft].pdf", EntryKey: "uuid/source [draft].pdf"}}, + `answer + +Sources: +[1] [source \[draft\].pdf](/api/agents/collections/agent%20one/entries-raw/uuid/source%20%5Bdraft%5D.pdf)`, + ), +) + var _ = Describe("ExecuteChatWithLLM", func() { var ( ctx context.Context @@ -184,6 +254,53 @@ var _ = Describe("ExecuteChatWithLLM", func() { }) }) + Context("knowledge base citations", func() { + It("appends KB sources to the returned response and callback message", func() { + mux := http.NewServeMux() + mux.HandleFunc("/api/agents/collections/kb-agent/search", func(w http.ResponseWriter, r *http.Request) { + Expect(r.URL.Query().Get("user_id")).To(Equal("user-1")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "results": [ + { + "content": "KB content", + "id": "result-1", + "similarity": 0.99, + "metadata": { + "file_name": "new feature.pdf", + "source": "uuid/new feature.pdf" + } + } + ], + "count": 1 + }`)) + }) + server := httptest.NewServer(mux) + defer server.Close() + + var msgContent string + cb.OnMessage = func(sender, content, messageID string) { + msgContent = content + } + + llm := &mockLLM{response: "agent reply"} + cfg := &AgentConfig{ + Name: "kb-agent", + Model: "test-model", + EnableKnowledgeBase: true, + KBMode: KBModeAutoSearch, + } + + result, err := ExecuteChatWithLLM(ctx, llm, cfg, "hello", cb, ExecuteChatOpts{ + APIURL: server.URL, + UserID: "user-1", + }) + Expect(err).ToNot(HaveOccurred()) + Expect(result).To(Equal("agent reply\n\nSources:\n[1] [new feature.pdf](/api/agents/collections/kb-agent/entries-raw/uuid/new%20feature.pdf?user_id=user-1)")) + Expect(msgContent).To(Equal(result)) + }) + }) + Context("context cancellation", func() { It("returns an error when context is already cancelled", func() { cancelledCtx, cancel := context.WithCancel(ctx) From 04f9608bafafdf2f6b069a335fc899cde1867c81 Mon Sep 17 00:00:00 2001 From: Ching Kao <0980124jim@gmail.com> Date: Tue, 9 Jun 2026 12:17:06 +0800 Subject: [PATCH 3/4] Collect KB citations from tool searches Signed-off-by: Pete Chen --- core/services/agents/citations.go | 127 ++++++++++++++++++++++++++ core/services/agents/executor.go | 105 +-------------------- core/services/agents/executor_test.go | 127 +++++++++++++++++++++++++- core/services/agents/knowledge.go | 33 +++++-- 4 files changed, 280 insertions(+), 112 deletions(-) create mode 100644 core/services/agents/citations.go diff --git a/core/services/agents/citations.go b/core/services/agents/citations.go new file mode 100644 index 000000000000..65c96d75c5b7 --- /dev/null +++ b/core/services/agents/citations.go @@ -0,0 +1,127 @@ +package agents + +import ( + "fmt" + "net/url" + "strings" + "sync" +) + +type kbCitationList struct { + mu sync.Mutex + citations []KBCitation +} + +func (l *kbCitationList) AddKBCitations(citations []KBCitation) { + if len(citations) == 0 { + return + } + l.mu.Lock() + defer l.mu.Unlock() + l.citations = append(l.citations, citations...) +} + +func (l *kbCitationList) Citations() []KBCitation { + l.mu.Lock() + defer l.mu.Unlock() + out := make([]KBCitation, len(l.citations)) + copy(out, l.citations) + return out +} + +// AppendKBCitations appends a markdown Sources block for KB citations. +func AppendKBCitations(response, collection, userID string, citations []KBCitation) string { + if strings.TrimSpace(response) == "" || len(citations) == 0 { + return response + } + + var lines []string + seen := make(map[string]struct{}) + for _, citation := range citations { + key := strings.TrimSpace(citation.EntryKey) + if key == "" { + key = strings.TrimSpace(citation.FileName) + } + if key == "" { + continue + } + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + + displayName := kbCitationDisplayName(citation) + if displayName == "" { + continue + } + + sourceURL := kbCitationRawFileURL(collection, citation.EntryKey, userID) + number := len(lines) + 1 + if sourceURL == "" { + lines = append(lines, fmt.Sprintf("[%d] %s", number, displayName)) + continue + } + lines = append(lines, fmt.Sprintf("[%d] [%s](%s)", number, escapeMarkdownLinkText(displayName), sourceURL)) + } + + if len(lines) == 0 { + return response + } + + var sb strings.Builder + sb.WriteString(strings.TrimRight(response, "\n")) + sb.WriteString("\n\nSources:\n") + for _, line := range lines { + sb.WriteString(line) + sb.WriteString("\n") + } + return strings.TrimRight(sb.String(), "\n") +} + +func kbCitationDisplayName(citation KBCitation) string { + if fileName := strings.TrimSpace(citation.FileName); fileName != "" { + return fileName + } + + segments := strings.Split(strings.Trim(strings.TrimSpace(citation.EntryKey), "/"), "/") + for i := len(segments) - 1; i >= 0; i-- { + if segment := strings.TrimSpace(segments[i]); segment != "" { + return segment + } + } + return "" +} + +func kbCitationRawFileURL(collection, entryKey, userID string) string { + collection = strings.TrimSpace(collection) + entryKey = strings.Trim(strings.TrimSpace(entryKey), "/") + if collection == "" || entryKey == "" { + return "" + } + + var escapedEntrySegments []string + for _, segment := range strings.Split(entryKey, "/") { + if segment == "" { + continue + } + escapedEntrySegments = append(escapedEntrySegments, url.PathEscape(segment)) + } + if len(escapedEntrySegments) == 0 { + return "" + } + + sourceURL := "/api/agents/collections/" + url.PathEscape(collection) + "/entries-raw/" + strings.Join(escapedEntrySegments, "/") + if userID != "" { + query := url.Values{} + query.Set("user_id", userID) + sourceURL += "?" + query.Encode() + } + return sourceURL +} + +func escapeMarkdownLinkText(text string) string { + text = strings.ReplaceAll(text, `\`, `\\`) + text = strings.ReplaceAll(text, "[", `\[`) + text = strings.ReplaceAll(text, "]", `\]`) + return text +} diff --git a/core/services/agents/executor.go b/core/services/agents/executor.go index 3d23da0c81bf..2787aeabc1af 100644 --- a/core/services/agents/executor.go +++ b/core/services/agents/executor.go @@ -5,7 +5,6 @@ import ( "context" "encoding/json" "fmt" - "net/url" "strings" "time" @@ -168,12 +167,12 @@ func ExecuteChatWithLLM(ctx context.Context, llm cogito.LLM, cfg *AgentConfig, m } } - var kbCitations []KBCitation + kbCitations := &kbCitationList{} if cfg.EnableKnowledgeBase && (kbMode == KBModeAutoSearch || kbMode == KBModeBoth) { kbResult := KBAutoSearchPrompt(ctx, effectiveURL, effectiveKey, cfg.Name, message, cfg.KnowledgeBaseResults, userID) if kbResult.Prompt != "" { fragment = fragment.AddMessage(cogito.SystemMessageRole, kbResult.Prompt) - kbCitations = kbResult.Citations + kbCitations.AddKBCitations(kbResult.Citations) } } @@ -200,7 +199,7 @@ func ExecuteChatWithLLM(ctx context.Context, llm cogito.LLM, cfg *AgentConfig, m } cogitoOpts = append(cogitoOpts, cogito.WithTools( cogito.NewToolDefinition( - KBSearchMemoryTool{APIURL: effectiveURL, APIKey: effectiveKey, Collection: cfg.Name, MaxResults: kbResults, UserID: userID}, + KBSearchMemoryTool{APIURL: effectiveURL, APIKey: effectiveKey, Collection: cfg.Name, MaxResults: kbResults, UserID: userID, CitationCollector: kbCitations}, KBSearchMemoryArgs{}, "search_memory", "Search the knowledge base for relevant information", @@ -340,7 +339,7 @@ func ExecuteChatWithLLM(ctx context.Context, llm cogito.LLM, cfg *AgentConfig, m response = stripThinkingTags(response) } responseForMemory := response - response = appendKBCitations(response, cfg.Name, userID, kbCitations) + response = AppendKBCitations(response, cfg.Name, userID, kbCitations.Citations()) // Save conversation to KB when long-term memory is enabled. // Use a detached context: the parent ctx may be cancelled (e.g. in distributed @@ -385,102 +384,6 @@ func isInternalCogitoTool(name string) bool { return false } -func appendKBCitations(response, collection, userID string, citations []KBCitation) string { - if strings.TrimSpace(response) == "" || len(citations) == 0 { - return response - } - - var lines []string - seen := make(map[string]struct{}) - for _, citation := range citations { - key := strings.TrimSpace(citation.EntryKey) - if key == "" { - key = strings.TrimSpace(citation.FileName) - } - if key == "" { - continue - } - if _, ok := seen[key]; ok { - continue - } - seen[key] = struct{}{} - - displayName := kbCitationDisplayName(citation) - if displayName == "" { - continue - } - - sourceURL := kbCitationRawFileURL(collection, citation.EntryKey, userID) - number := len(lines) + 1 - if sourceURL == "" { - lines = append(lines, fmt.Sprintf("[%d] %s", number, displayName)) - continue - } - lines = append(lines, fmt.Sprintf("[%d] [%s](%s)", number, escapeMarkdownLinkText(displayName), sourceURL)) - } - - if len(lines) == 0 { - return response - } - - var sb strings.Builder - sb.WriteString(strings.TrimRight(response, "\n")) - sb.WriteString("\n\nSources:\n") - for _, line := range lines { - sb.WriteString(line) - sb.WriteString("\n") - } - return strings.TrimRight(sb.String(), "\n") -} - -func kbCitationDisplayName(citation KBCitation) string { - if fileName := strings.TrimSpace(citation.FileName); fileName != "" { - return fileName - } - - segments := strings.Split(strings.Trim(strings.TrimSpace(citation.EntryKey), "/"), "/") - for i := len(segments) - 1; i >= 0; i-- { - if segment := strings.TrimSpace(segments[i]); segment != "" { - return segment - } - } - return "" -} - -func kbCitationRawFileURL(collection, entryKey, userID string) string { - collection = strings.TrimSpace(collection) - entryKey = strings.Trim(strings.TrimSpace(entryKey), "/") - if collection == "" || entryKey == "" { - return "" - } - - var escapedEntrySegments []string - for _, segment := range strings.Split(entryKey, "/") { - if segment == "" { - continue - } - escapedEntrySegments = append(escapedEntrySegments, url.PathEscape(segment)) - } - if len(escapedEntrySegments) == 0 { - return "" - } - - sourceURL := "/api/agents/collections/" + url.PathEscape(collection) + "/entries-raw/" + strings.Join(escapedEntrySegments, "/") - if userID != "" { - query := url.Values{} - query.Set("user_id", userID) - sourceURL += "?" + query.Encode() - } - return sourceURL -} - -func escapeMarkdownLinkText(text string) string { - text = strings.ReplaceAll(text, `\`, `\\`) - text = strings.ReplaceAll(text, "[", `\[`) - text = strings.ReplaceAll(text, "]", `\]`) - return text -} - // stripThinkingTags removes ... blocks from content. func stripThinkingTags(content string) string { // Simple implementation — remove content between and diff --git a/core/services/agents/executor_test.go b/core/services/agents/executor_test.go index 5b86005fc3da..9b2001ccaf10 100644 --- a/core/services/agents/executor_test.go +++ b/core/services/agents/executor_test.go @@ -38,6 +38,34 @@ func (m *mockLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompl }, cogito.LLMUsage{}, nil } +type toolCallingMockLLM struct { + createResponses []openai.ChatCompletionResponse + askResponse string + callCount atomic.Int32 +} + +func (m *toolCallingMockLLM) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) { + m.callCount.Add(1) + return f.AddMessage(cogito.AssistantMessageRole, m.askResponse), nil +} + +func (m *toolCallingMockLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) { + idx := int(m.callCount.Add(1)) - 1 + if idx >= len(m.createResponses) { + return cogito.LLMReply{ + ChatCompletionResponse: openai.ChatCompletionResponse{ + Choices: []openai.ChatCompletionChoice{{ + Message: openai.ChatCompletionMessage{ + Role: "assistant", + Content: "No more tools needed.", + }, + }}, + }, + }, cogito.LLMUsage{}, nil + } + return cogito.LLMReply{ChatCompletionResponse: m.createResponses[idx]}, cogito.LLMUsage{}, nil +} + // statusCollector records status callbacks in a thread-safe way. type statusCollector struct { mu sync.Mutex @@ -77,7 +105,7 @@ var _ = DescribeTable("stripThinkingTags", var _ = DescribeTable("appendKBCitations", func(response, collection, userID string, citations []KBCitation, want string) { - Expect(appendKBCitations(response, collection, userID, citations)).To(Equal(want)) + Expect(AppendKBCitations(response, collection, userID, citations)).To(Equal(want)) }, Entry("leaves responses without citations unchanged", "answer", @@ -299,6 +327,103 @@ var _ = Describe("ExecuteChatWithLLM", func() { Expect(result).To(Equal("agent reply\n\nSources:\n[1] [new feature.pdf](/api/agents/collections/kb-agent/entries-raw/uuid/new%20feature.pdf?user_id=user-1)")) Expect(msgContent).To(Equal(result)) }) + + It("collects citations from the search_memory tool", func() { + mux := http.NewServeMux() + mux.HandleFunc("/api/agents/collections/kb-agent/search", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "results": [ + { + "content": "Tool KB content", + "id": "result-1", + "similarity": 0.99, + "metadata": { + "file_name": "tool source.pdf", + "source": "uuid/tool source.pdf" + } + } + ], + "count": 1 + }`)) + }) + server := httptest.NewServer(mux) + defer server.Close() + + collector := &kbCitationList{} + tool := KBSearchMemoryTool{ + APIURL: server.URL, + Collection: "kb-agent", + CitationCollector: collector, + } + + result, _, err := tool.Run(KBSearchMemoryArgs{Query: "hello"}) + Expect(err).ToNot(HaveOccurred()) + Expect(result).To(ContainSubstring("Tool KB content")) + Expect(collector.Citations()).To(Equal([]KBCitation{{FileName: "tool source.pdf", EntryKey: "uuid/tool source.pdf"}})) + }) + + It("appends KB sources found through tools-only search_memory calls", func() { + mux := http.NewServeMux() + mux.HandleFunc("/api/agents/collections/kb-agent/search", func(w http.ResponseWriter, r *http.Request) { + Expect(r.URL.Query().Get("user_id")).To(Equal("user-1")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "results": [ + { + "content": "Tool KB content", + "id": "result-1", + "similarity": 0.99, + "metadata": { + "file_name": "tool source.pdf", + "source": "uuid/tool source.pdf" + } + } + ], + "count": 1 + }`)) + }) + server := httptest.NewServer(mux) + defer server.Close() + + llm := &toolCallingMockLLM{ + askResponse: "agent reply from tool context", + createResponses: []openai.ChatCompletionResponse{ + { + Choices: []openai.ChatCompletionChoice{ + { + Message: openai.ChatCompletionMessage{ + Role: "assistant", + ToolCalls: []openai.ToolCall{ + { + ID: "call-1", + Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{ + Name: "search_memory", + Arguments: `{"query":"hello"}`, + }, + }, + }, + }, + }, + }, + }, + }, + } + cfg := &AgentConfig{ + Name: "kb-agent", + Model: "test-model", + EnableKnowledgeBase: true, + KBMode: KBModeTools, + } + + result, err := ExecuteChatWithLLM(ctx, llm, cfg, "hello", cb, ExecuteChatOpts{ + APIURL: server.URL, + UserID: "user-1", + }) + Expect(err).ToNot(HaveOccurred()) + Expect(result).To(Equal("agent reply from tool context\n\nSources:\n[1] [tool source.pdf](/api/agents/collections/kb-agent/entries-raw/uuid/tool%20source.pdf?user_id=user-1)")) + }) }) Context("context cancellation", func() { diff --git a/core/services/agents/knowledge.go b/core/services/agents/knowledge.go index 54ae57be5167..ae87f560d917 100644 --- a/core/services/agents/knowledge.go +++ b/core/services/agents/knowledge.go @@ -8,6 +8,7 @@ import ( "io" "mime/multipart" "net/http" + "net/url" "strings" "time" @@ -59,6 +60,11 @@ type KBSearchContext struct { Citations []KBCitation `json:"citations"` } +// KBCitationCollector receives source citations found during KB searches. +type KBCitationCollector interface { + AddKBCitations([]KBCitation) +} + // KBAutoSearchPrompt queries the knowledge base with the user's message and // returns a KBSearchContext: a system prompt block with the relevant results // plus the de-duplicated source citations those results came from. @@ -72,10 +78,11 @@ func KBAutoSearchPrompt(ctx context.Context, apiURL, apiKey, collection, query s maxResults = 5 } - // Call LocalAI's collection search API - searchURL := strings.TrimRight(apiURL, "/") + "/api/agents/collections/" + collection + "/search" + searchURL := strings.TrimRight(apiURL, "/") + "/api/agents/collections/" + url.PathEscape(collection) + "/search" if userID != "" { - searchURL += "?user_id=" + userID + query := url.Values{} + query.Set("user_id", userID) + searchURL += "?" + query.Encode() } reqBody, _ := json.Marshal(map[string]any{ "query": query, @@ -168,11 +175,12 @@ type KBSearchMemoryArgs struct { // KBSearchMemoryTool implements the search_memory MCP tool. type KBSearchMemoryTool struct { - APIURL string - APIKey string - Collection string - MaxResults int - UserID string + APIURL string + APIKey string + Collection string + MaxResults int + UserID string + CitationCollector KBCitationCollector } func (t KBSearchMemoryTool) Run(args KBSearchMemoryArgs) (string, any, error) { @@ -182,6 +190,9 @@ func (t KBSearchMemoryTool) Run(args KBSearchMemoryArgs) (string, any, error) { if result.Prompt == "" { return "No results found.", nil, nil } + if t.CitationCollector != nil { + t.CitationCollector.AddKBCitations(result.Citations) + } return result.Prompt, nil, nil } @@ -214,9 +225,11 @@ func (t KBAddMemoryTool) Run(args KBAddMemoryArgs) (string, any, error) { // KBStoreContent uploads text content to a collection via the multipart upload API. func KBStoreContent(ctx context.Context, apiURL, apiKey, collection, content, userID string) error { - uploadURL := strings.TrimRight(apiURL, "/") + "/api/agents/collections/" + collection + "/upload" + uploadURL := strings.TrimRight(apiURL, "/") + "/api/agents/collections/" + url.PathEscape(collection) + "/upload" if userID != "" { - uploadURL += "?user_id=" + userID + query := url.Values{} + query.Set("user_id", userID) + uploadURL += "?" + query.Encode() } // Build multipart form with the text content as a file From a621190f405a4493cc79d94efb11de85139a026d Mon Sep 17 00:00:00 2001 From: Ching Kao <0980124jim@gmail.com> Date: Tue, 9 Jun 2026 14:51:52 +0800 Subject: [PATCH 4/4] fix(agents): append KB sources in local chats Apply the shared KB citation post-processing to standalone LocalAGI chat responses so the React agent chat receives the same clickable Sources block as the native executor path. Also fix the run target to use the current cmd/local-ai entrypoint. Assisted-by: Codex:gpt-5 Signed-off-by: Pete Chen --- Makefile | 2 +- core/services/agentpool/agent_pool.go | 76 ++++++++++++++++++++++++++- 2 files changed, 76 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 9503bf7b483b..cafcdd44a692 100644 --- a/Makefile +++ b/Makefile @@ -180,7 +180,7 @@ osx-signed: build ## Run run: ## run local-ai - CGO_LDFLAGS="$(CGO_LDFLAGS)" $(GOCMD) run ./ + CGO_LDFLAGS="$(CGO_LDFLAGS)" $(GOCMD) run ./cmd/local-ai prepare-test: protogen-go build-mock-backend diff --git a/core/services/agentpool/agent_pool.go b/core/services/agentpool/agent_pool.go index 9216add9b3c5..8744b8a35d9e 100644 --- a/core/services/agentpool/agent_pool.go +++ b/core/services/agentpool/agent_pool.go @@ -466,10 +466,11 @@ func (s *AgentPoolService) Chat(name, message string) (string, error) { s.collectAndCopyMetadata(metadata, chatUserID) } + content := s.appendLocalAGIKBCitations(response.Response, name, message, response.State) msg := map[string]any{ "id": messageID + "-agent", "sender": "agent", - "content": response.Response, + "content": content, "timestamp": time.Now().Format(time.RFC3339), } if len(metadata) > 0 { @@ -489,6 +490,79 @@ func (s *AgentPoolService) Chat(name, message string) (string, error) { return messageID, nil } +func (s *AgentPoolService) appendLocalAGIKBCitations(response, agentKey, message string, states []coreTypes.ActionState) string { + if strings.TrimSpace(response) == "" { + return response + } + + userID, collection := splitAgentKey(agentKey) + cfg := s.localAGI.pool.GetConfig(agentKey) + if cfg == nil || !cfg.EnableKnowledgeBase { + return response + } + + citations := kbCitationsFromActionStates(states) + if len(citations) == 0 && cfg.KBAutoSearch { + maxResults := cfg.KnowledgeBaseResults + if maxResults <= 0 { + maxResults = 5 + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + kbResult := agents.KBAutoSearchPrompt(ctx, s.apiURL, s.apiKey, collection, message, maxResults, userID) + citations = kbResult.Citations + } + + return agents.AppendKBCitations(response, collection, userID, citations) +} + +func splitAgentKey(agentKey string) (userID, name string) { + if uid, n, ok := strings.Cut(agentKey, ":"); ok { + return uid, n + } + return "", agentKey +} + +func kbCitationsFromActionStates(states []coreTypes.ActionState) []agents.KBCitation { + var citations []agents.KBCitation + for _, state := range states { + citations = append(citations, kbCitationsFromMetadata(state.Metadata)...) + } + return citations +} + +func kbCitationsFromMetadata(metadata map[string]any) []agents.KBCitation { + if len(metadata) == 0 { + return nil + } + + fileName := metadata["file_name"] + source := metadata["source"] + if fileName == nil && source == nil { + return nil + } + + citation := agents.KBCitation{ + FileName: metadataString(fileName), + EntryKey: metadataString(source), + } + if citation.FileName == "" && citation.EntryKey == "" { + return nil + } + return []agents.KBCitation{citation} +} + +func metadataString(value any) string { + switch v := value.(type) { + case string: + return v + case fmt.Stringer: + return v.String() + default: + return "" + } +} + // userOutputsDir returns the per-user outputs directory, creating it if needed. // If userID is empty, falls back to the shared outputs directory. func (s *AgentPoolService) userOutputsDir(userID string) string {