feat(agents): surface KB source citations in RAG responses (#10228)

* dev knowledge.go structure Signed-off-by: Pete Chen <petechentw@gmail.com> * 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 <petechentw@gmail.com> * Collect KB citations from tool searches Signed-off-by: Pete Chen <petechentw@gmail.com> * 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 <petechentw@gmail.com> --------- Signed-off-by: Pete Chen <petechentw@gmail.com> Co-authored-by: shihyunhuang <shihyunhuang88@gmail.com> Co-authored-by: TLoE419 <tloemizuchizu@gmail.com> Co-authored-by: Ching Kao <0980124jim@gmail.com>
2026-06-10 17:56:49 -04:00 · 2026-06-09 07:32:56 -07:00
parent e1ec03d33f
commit d2e6b93369
6 changed files with 556 additions and 38 deletions
--- a/2
+++ b/2
@@ -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

--- 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 {
--- a/core/services/agents/citations.go
+++ 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
+}
--- a/core/services/agents/executor.go
+++ b/core/services/agents/executor.go
@@ -167,10 +167,12 @@ func ExecuteChatWithLLM(ctx context.Context, llm cogito.LLM, cfg *AgentConfig, m
 		}
 	}

+	kbCitations := &kbCitationList{}
 	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)
+			kbCitations.AddKBCitations(kbResult.Citations)
 		}
 	}

@@ -197,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",
@@ -336,6 +338,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.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
@@ -344,7 +348,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)
 		}()
 	}

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

@@ -36,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
@@ -73,6 +103,74 @@ var _ = DescribeTable("stripThinkingTags",
 	Entry("adjacent tag pairs", "<thinking>a</thinking><thinking>b</thinking>", ""),
 )

+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 +282,150 @@ 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))
+		})
+
+		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() {
 		It("returns an error when context is already cancelled", func() {
 			cancelledCtx, cancel := context.WithCancel(ctx)
--- 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"

@@ -17,10 +18,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 ("<uuid>/<filename>") 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,22 +41,48 @@ 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 ("<uuid>/<filename>"), 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"`
+}
+
+// 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.
 // 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 {
 		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,
@@ -56,7 +92,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 +102,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.
@@ -110,21 +175,25 @@ 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) {
 	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
+	if t.CitationCollector != nil {
+		t.CitationCollector.AddKBCitations(result.Citations)
+	}
+	return result.Prompt, nil, nil
 }

 // KBAddMemoryArgs defines the arguments for the add_memory tool.
@@ -156,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