feat(backend): improve parent-child chunk merge and imageinfo matching

Scope image info to matched text children during parent resolution and add range/URL filtering utilities with unit tests.
2026-06-16 00:49:16 +08:00
parent 8dbc20b9b3
commit 2906c6656e
7 changed files with 457 additions and 128 deletions
@@ -304,5 +304,5 @@ func getEnrichedPassageForChat(ctx context.Context, result *types.SearchResult)

 // enrichContentWithImageInfo delegates to the shared searchutil implementation.
 func enrichContentWithImageInfo(_ context.Context, content string, imageInfoJSON string) string {
-	return searchutil.EnrichContentWithImageInfo(content, imageInfoJSON)
+	return searchutil.EnrichContentWithImageInfoForChat(content, imageInfoJSON)
 }
@@ -2,9 +2,9 @@ package chatpipeline

 import (
 	"context"
-	"encoding/json"
 	"sort"

+	"github.com/Tencent/WeKnora/internal/searchutil"
 	"github.com/Tencent/WeKnora/internal/types"
 	"github.com/Tencent/WeKnora/internal/types/interfaces"
 )
@@ -39,7 +39,7 @@ func (p *PluginMerge) ActivationEvents() []types.EventType {
 //  5. Group by knowledge source + chunk type, merge overlapping ranges
 //  6. Populate FAQ answers
 //  7. Expand short contexts with neighboring chunks
-//  7.5. Re-merge overlapping ranges introduced by expansion
+//     7.5. Re-merge overlapping ranges introduced by expansion
 //  8. Final deduplication (ID + signature + partial content overlap)
 func (p *PluginMerge) OnEvent(ctx context.Context,
 	eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
@@ -207,9 +207,11 @@ func (p *PluginMerge) groupAndMergeOverlapping(ctx context.Context, results []*t
 	return mergedChunks
 }

-// resolveParentChunks replaces child chunk content with parent chunk content
-// for results that have ParentChunkID set. This provides fuller context
-// for small child chunks used in parent-child chunking strategy.
+// resolveParentChunks scopes parent-child retrieval results to the matched
+// child window instead of expanding to the full parent_text block. Text
+// children keep their own content; image children resolve to the markdown
+// slice of the parent_text that covers their text parent. ImageInfo is
+// collected only for the matched text child, not all siblings under parent.
 func (p *PluginMerge) resolveParentChunks(
 	ctx context.Context,
 	chatManage *types.ChatManage,
@@ -294,14 +296,12 @@ func (p *PluginMerge) resolveParentChunks(
 		}
 	}

-	// Collect merged ImageInfo for each parent by fetching ALL sibling
-	// child chunks. Individual child chunks only carry ImageInfo for images
-	// within their own range, but the parent content spans all children.
-	imageInfoIDs := ids
-	if len(grandparentIDs) > 0 {
-		imageInfoIDs = append(append([]string(nil), ids...), grandparentIDs...)
+	// Batch-fetch image_info scoped to matched text children only.
+	textChildIDs := collectScopedTextChildIDs(results, parentMap)
+	var scopedImageInfo map[string]string
+	if len(textChildIDs) > 0 {
+		scopedImageInfo = searchutil.CollectImageInfoByChunkIDs(ctx, p.chunkRepo, tenantID, textChildIDs)
 	}
-	parentImageInfoMap := p.collectParentImageInfo(ctx, tenantID, imageInfoIDs)

 	for _, r := range results {
 		if r.ParentChunkID == "" {
@@ -310,56 +310,74 @@ func (p *PluginMerge) resolveParentChunks(

 		switch r.ChunkType {
 		case string(types.ChunkTypeText):
-			// text → parent_text resolution (parent-child chunking strategy).
-			// Summary chunks also carry a ParentChunkID that points to their
-			// source chunk, but that is a different semantic — replacing
-			// summary content with its source would degrade quality.
+			// text → parent_text: expand to full parent for surrounding context
+			// (the core parent-child value). Scope ImageInfo to this child only so
+			// image-heavy parents do not inject every sibling page's OCR/Caption.
 			parent, ok := parentMap[r.ParentChunkID]
 			if !ok || parent.Content == "" || parent.ChunkType != types.ChunkTypeParentText {
 				continue
 			}
+			matchStart, matchEnd := r.StartAt, r.EndAt
 			pipelineInfo(ctx, "Merge", "parent_resolve", map[string]interface{}{
 				"child_id":   r.ID,
 				"parent_id":  r.ParentChunkID,
 				"child_len":  runeLen(r.Content),
 				"parent_len": runeLen(parent.Content),
+				"scoped_img": true,
 			})
-			r.Content = parent.Content
+			r.Content = searchutil.PruneMarkdownImagesOutsideRange(
+				parent.Content, parent.StartAt, matchStart, matchEnd,
+			)
 			r.StartAt = parent.StartAt
 			r.EndAt = parent.EndAt
-			if mergedImageInfo, ok := parentImageInfoMap[r.ParentChunkID]; ok && mergedImageInfo != "" {
-				r.ImageInfo = mergedImageInfo
+			assignScopedImageInfo(r, scopedImageInfo, r.ID)
+			if r.ImageInfo != "" {
+				r.ImageInfo = searchutil.FilterImageInfoByMatchRange(
+					parent.Content, parent.StartAt, matchStart, matchEnd, r.ImageInfo,
+				)
 			}
 			if !containsID(r.SubChunkID, r.ID) {
 				r.SubChunkID = append(r.SubChunkID, r.ID)
 			}

 		case string(types.ChunkTypeImageOCR), string(types.ChunkTypeImageCaption):
-			// image_ocr/image_caption → text parent → optional parent_text grandparent.
-			// Replace content with parent text for surrounding context.
-			parent, ok := parentMap[r.ParentChunkID]
-			if !ok || parent.Content == "" {
+			textParent, ok := parentMap[r.ParentChunkID]
+			if !ok || textParent.Content == "" || textParent.ChunkType != types.ChunkTypeText {
 				continue
 			}
-			resolvedParent := parent
-			// If parent text uses parent-child chunking, resolve one more level
-			if parent.ChunkType == types.ChunkTypeText && parent.ParentChunkID != "" {
-				if gp, gpOK := parentMap[parent.ParentChunkID]; gpOK && gp.ChunkType == types.ChunkTypeParentText && gp.Content != "" {
-					resolvedParent = gp
+			hitImageInfo := r.ImageInfo
+			contentSource := textParent
+			if textParent.ParentChunkID != "" {
+				if gp, gpOK := parentMap[textParent.ParentChunkID]; gpOK &&
+					gp.ChunkType == types.ChunkTypeParentText && gp.Content != "" {
+					contentSource = gp
 				}
 			}
+			matchStart := textParent.StartAt
+			matchEnd := textParent.EndAt
+			sliced := searchutil.SliceContentByDocumentRange(
+				contentSource.Content, contentSource.StartAt, matchStart, matchEnd,
+			)
+			if sliced == "" {
+				sliced = textParent.Content
+				matchStart = textParent.StartAt
+				matchEnd = textParent.EndAt
+			}
 			pipelineInfo(ctx, "Merge", "image_parent_resolve", map[string]interface{}{
-				"child_id":    r.ID,
-				"child_type":  r.ChunkType,
-				"resolved_id": resolvedParent.ID,
-				"child_len":   runeLen(r.Content),
-				"parent_len":  runeLen(resolvedParent.Content),
+				"child_id":   r.ID,
+				"child_type": r.ChunkType,
+				"text_id":    textParent.ID,
+				"parent_id":  contentSource.ID,
+				"match_len":  runeLen(sliced),
+				"parent_len": runeLen(contentSource.Content),
+				"scoped":     true,
 			})
-			r.Content = resolvedParent.Content
-			r.StartAt = resolvedParent.StartAt
-			r.EndAt = resolvedParent.EndAt
-			if mergedInfo, ok := parentImageInfoMap[resolvedParent.ID]; ok && mergedInfo != "" {
-				r.ImageInfo = mergedInfo
+			r.Content = sliced
+			r.StartAt = matchStart
+			r.EndAt = matchEnd
+			assignScopedImageInfo(r, scopedImageInfo, textParent.ID)
+			if r.ImageInfo == "" && hitImageInfo != "" {
+				r.ImageInfo = searchutil.FilterImageInfoByContentURLs(r.Content, hitImageInfo)
 			}
 			if !containsID(r.SubChunkID, r.ID) {
 				r.SubChunkID = append(r.SubChunkID, r.ID)
@@ -370,97 +388,50 @@ func (p *PluginMerge) resolveParentChunks(
 	return results
 }

-// collectParentImageInfo batch-fetches image info for the given parent_text
-// chunk IDs using a two-level query:
-//
-//	Level 1: parent_text → text children
-//	Level 2: text children → image_ocr / image_caption grandchildren (carry image_info)
-func (p *PluginMerge) collectParentImageInfo(
-	ctx context.Context,
-	tenantID uint64,
-	parentIDs []string,
-) map[string]string {
-	// Level 1: get direct children of parent_text chunks
-	children, err := p.chunkRepo.ListChunksByParentIDs(ctx, tenantID, parentIDs)
-	if err != nil {
-		pipelineWarn(ctx, "Merge", "parent_imageinfo_fetch_failed", map[string]interface{}{
-			"parent_cnt": len(parentIDs),
-			"error":      err.Error(),
-		})
-		return nil
-	}
-
-	type agg struct {
-		infos    []types.ImageInfo
-		seenURLs map[string]bool
-	}
-	aggMap := make(map[string]*agg)
-	addInfo := func(targetID string, chunk *types.Chunk) {
-		if chunk.ImageInfo == "" {
-			return
-		}
-		var infos []types.ImageInfo
-		if err := json.Unmarshal([]byte(chunk.ImageInfo), &infos); err != nil || len(infos) == 0 {
-			return
-		}
-		a, ok := aggMap[targetID]
-		if !ok {
-			a = &agg{seenURLs: make(map[string]bool)}
-			aggMap[targetID] = a
-		}
-		for _, info := range infos {
-			key := info.URL
-			if key == "" {
-				key = info.OriginalURL
-			}
-			if key != "" && !a.seenURLs[key] {
-				a.seenURLs[key] = true
-				a.infos = append(a.infos, info)
-			}
-		}
-	}
-
-	var textChildIDs []string
-	textToParent := make(map[string]string, len(children))
-	for _, child := range children {
-		switch child.ChunkType {
-		case types.ChunkTypeImageOCR, types.ChunkTypeImageCaption:
-			addInfo(child.ParentChunkID, child)
-		case types.ChunkTypeText:
-			textChildIDs = append(textChildIDs, child.ID)
-			textToParent[child.ID] = child.ParentChunkID
-		}
-	}
-
-	// Level 2: text children → image grandchildren
-	if len(textChildIDs) > 0 {
-		grandChildren, err := p.chunkRepo.ListChunksByParentIDs(ctx, tenantID, textChildIDs)
-		if err != nil {
-			pipelineWarn(ctx, "Merge", "parent_imageinfo_l2_fetch_failed", map[string]interface{}{
-				"text_cnt": len(textChildIDs),
-				"error":    err.Error(),
-			})
-		} else {
-			for _, gc := range grandChildren {
-				if gc.ChunkType != types.ChunkTypeImageOCR && gc.ChunkType != types.ChunkTypeImageCaption {
-					continue
-				}
-				if parentTextID, ok := textToParent[gc.ParentChunkID]; ok {
-					addInfo(parentTextID, gc)
-				}
-			}
-		}
-	}
-
-	result := make(map[string]string, len(aggMap))
-	for id, a := range aggMap {
-		if len(a.infos) == 0 {
+// collectScopedTextChildIDs returns text chunk IDs whose image_info should be
+// loaded for parent-child merge scoping.
+func collectScopedTextChildIDs(
+	results []*types.SearchResult,
+	parentMap map[string]*types.Chunk,
+) []string {
+	seen := make(map[string]struct{})
+	var ids []string
+	for _, r := range results {
+		if r.ParentChunkID == "" {
 			continue
 		}
-		data, err := json.Marshal(a.infos)
-		if err == nil {
-			result[id] = string(data)
+		switch r.ChunkType {
+		case string(types.ChunkTypeText):
+			parent := parentMap[r.ParentChunkID]
+			if parent == nil || parent.ChunkType != types.ChunkTypeParentText {
+				continue
+			}
+			if _, ok := seen[r.ID]; ok {
+				continue
+			}
+			seen[r.ID] = struct{}{}
+			ids = append(ids, r.ID)
+		case string(types.ChunkTypeImageOCR), string(types.ChunkTypeImageCaption):
+			if _, ok := seen[r.ParentChunkID]; ok {
+				continue
+			}
+			seen[r.ParentChunkID] = struct{}{}
+			ids = append(ids, r.ParentChunkID)
 		}
 	}
-	return result
+	return ids
+}
+
+// assignScopedImageInfo sets ImageInfo from the per-text-child map, falling
+// back to URLs present in the result content.
+func assignScopedImageInfo(r *types.SearchResult, scoped map[string]string, textChildID string) {
+	if scoped != nil {
+		if info, ok := scoped[textChildID]; ok && info != "" {
+			r.ImageInfo = info
+			return
+		}
+	}
+	if r.ImageInfo != "" {
+		r.ImageInfo = searchutil.FilterImageInfoByContentURLs(r.Content, r.ImageInfo)
+	}
 }
@@ -0,0 +1,66 @@
+package chatpipeline
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/searchutil"
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+func TestCollectScopedTextChildIDs(t *testing.T) {
+	parentMap := map[string]*types.Chunk{
+		"parent-1": {ID: "parent-1", ChunkType: types.ChunkTypeParentText},
+		"text-x":   {ID: "text-x", ChunkType: types.ChunkTypeText},
+	}
+	results := []*types.SearchResult{
+		{ID: "text-1", ChunkType: string(types.ChunkTypeText), ParentChunkID: "parent-1"},
+		{ID: "img-1", ChunkType: string(types.ChunkTypeImageOCR), ParentChunkID: "text-2"},
+		{ID: "text-3", ChunkType: string(types.ChunkTypeText), ParentChunkID: "text-x"}, // not parent_text
+	}
+	ids := collectScopedTextChildIDs(results, parentMap)
+	if len(ids) != 2 {
+		t.Fatalf("ids: %v", ids)
+	}
+}
+
+func TestAssignScopedImageInfo_FiltersToContentURLs(t *testing.T) {
+	all, _ := json.Marshal([]types.ImageInfo{
+		{URL: "u1", OCRText: "one"},
+		{URL: "u2", OCRText: "two"},
+	})
+	r := &types.SearchResult{
+		Content:   "![p2](u2)",
+		ImageInfo: string(all),
+	}
+	assignScopedImageInfo(r, nil, "missing-child")
+	var infos []types.ImageInfo
+	if err := json.Unmarshal([]byte(r.ImageInfo), &infos); err != nil {
+		t.Fatal(err)
+	}
+	if len(infos) != 1 || infos[0].URL != "u2" {
+		t.Fatalf("filtered: %+v", infos)
+	}
+}
+
+func TestParentChildImageHit_WindowSliceAndFilter(t *testing.T) {
+	parentContent := "![p1](u1)\n\n![p2](u2)\n\n![p3](u3)"
+	textStart := len([]rune("![p1](u1)\n\n"))
+	textEnd := textStart + len([]rune("![p2](u2)"))
+	sliced := searchutil.SliceContentByDocumentRange(parentContent, 0, textStart, textEnd)
+	if sliced != "![p2](u2)" {
+		t.Fatalf("slice: got %q", sliced)
+	}
+
+	all, _ := json.Marshal([]types.ImageInfo{
+		{URL: "u1"}, {URL: "u2"}, {URL: "u3"},
+	})
+	filtered := searchutil.FilterImageInfoByContentURLs(sliced, string(all))
+	var infos []types.ImageInfo
+	if err := json.Unmarshal([]byte(filtered), &infos); err != nil {
+		t.Fatal(err)
+	}
+	if len(infos) != 1 || infos[0].URL != "u2" {
+		t.Fatalf("infos: %+v", infos)
+	}
+}
@@ -102,6 +102,9 @@ func buildStreamResponse(evt interfaces.StreamEvent, requestID string) *types.St
 	// Special handling for references event
 	if evt.Type == types.ResponseTypeReferences {
 		refsData := evt.Data["references"]
+		if refsData == nil {
+			return response
+		}
 		if refs, ok := refsData.(types.References); ok {
 			response.KnowledgeReferences = refs
 		} else if refs, ok := refsData.([]*types.SearchResult); ok {
@@ -260,6 +260,52 @@ func EnrichContentWithImageInfo(content string, imageInfoJSON string) string {
 	return content
 }

+// EnrichContentWithImageInfoForChat is like EnrichContentWithImageInfo but only
+// wraps Markdown images that have a matching image_info entry. This avoids
+// turning every parent thumbnail into an <image> block when only a few pages
+// were retrieved, and skips appending orphan image_info extras.
+func EnrichContentWithImageInfoForChat(content string, imageInfoJSON string) string {
+	var imageInfos []types.ImageInfo
+	if err := json.Unmarshal([]byte(imageInfoJSON), &imageInfos); err != nil {
+		return content
+	}
+	if len(imageInfos) == 0 {
+		return content
+	}
+
+	imageInfoMap := make(map[string]*types.ImageInfo)
+	for i := range imageInfos {
+		if imageInfos[i].URL != "" {
+			imageInfoMap[imageInfos[i].URL] = &imageInfos[i]
+		}
+		if imageInfos[i].OriginalURL != "" {
+			imageInfoMap[imageInfos[i].OriginalURL] = &imageInfos[i]
+		}
+	}
+
+	matches := MarkdownImageRegex.FindAllStringSubmatch(content, -1)
+	for _, match := range matches {
+		if len(match) < 3 {
+			continue
+		}
+		imgURL := match[2]
+		imgInfo, found := imageInfoMap[imgURL]
+		if !found || imgInfo == nil {
+			continue
+		}
+		inner := BuildImageInfoXML(imgInfo)
+		if inner == "" {
+			continue
+		}
+		var b strings.Builder
+		b.WriteString(fmt.Sprintf("<image url=\"%s\">\n", imgURL))
+		b.WriteString(inner)
+		b.WriteString("</image>")
+		content = strings.Replace(content, match[0], b.String(), 1)
+	}
+	return content
+}
+
 // BuildImageInfoXML returns XML-tagged caption / ocr for one image.
 func BuildImageInfoXML(img *types.ImageInfo) string {
 	var b strings.Builder
@@ -0,0 +1,148 @@
+package searchutil
+
+import (
+	"encoding/json"
+	"strings"
+	"unicode/utf8"
+
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+// SliceContentByDocumentRange extracts the substring of content whose
+// document-level rune offsets fall within [rangeStart, rangeEnd).
+// contentStartAt is the document offset where content begins.
+func SliceContentByDocumentRange(content string, contentStartAt, rangeStart, rangeEnd int) string {
+	if content == "" {
+		return content
+	}
+	relStart := rangeStart - contentStartAt
+	relEnd := rangeEnd - contentStartAt
+	runes := []rune(content)
+	if relStart < 0 {
+		relStart = 0
+	}
+	if relEnd > len(runes) {
+		relEnd = len(runes)
+	}
+	if relStart >= relEnd {
+		return ""
+	}
+	return string(runes[relStart:relEnd])
+}
+
+// ImageURLsInContent returns the set of image URLs referenced as Markdown
+// images in content.
+func ImageURLsInContent(content string) map[string]bool {
+	urls := make(map[string]bool)
+	for _, match := range MarkdownImageRegex.FindAllStringSubmatch(content, -1) {
+		if len(match) < 3 || match[2] == "" {
+			continue
+		}
+		urls[match[2]] = true
+	}
+	return urls
+}
+
+// FilterImageInfoByContentURLs keeps only image_info entries whose URL or
+// OriginalURL appears as a Markdown image in content. Returns empty string
+// when nothing matches or imageInfoJSON is invalid.
+func FilterImageInfoByContentURLs(content string, imageInfoJSON string) string {
+	if imageInfoJSON == "" {
+		return ""
+	}
+	var infos []types.ImageInfo
+	if err := json.Unmarshal([]byte(imageInfoJSON), &infos); err != nil || len(infos) == 0 {
+		return ""
+	}
+	urls := ImageURLsInContent(content)
+	if len(urls) == 0 {
+		return ""
+	}
+	filtered := make([]types.ImageInfo, 0, len(infos))
+	for _, info := range infos {
+		if urls[info.URL] || urls[info.OriginalURL] {
+			filtered = append(filtered, info)
+		}
+	}
+	return marshalImageInfos(filtered)
+}
+
+// FilterImageInfoByMatchRange keeps image_info entries whose Markdown image
+// reference falls within the document rune range [matchStart, matchEnd).
+// Used when parent content is expanded for context but multimodal enrichment
+// should only cover the retrieved child window.
+func FilterImageInfoByMatchRange(
+	parentContent string,
+	parentStartAt, matchStart, matchEnd int,
+	imageInfoJSON string,
+) string {
+	if imageInfoJSON == "" {
+		return ""
+	}
+	var infos []types.ImageInfo
+	if err := json.Unmarshal([]byte(imageInfoJSON), &infos); err != nil || len(infos) == 0 {
+		return ""
+	}
+	window := SliceContentByDocumentRange(parentContent, parentStartAt, matchStart, matchEnd)
+	urls := ImageURLsInContent(window)
+	if len(urls) == 0 {
+		return ""
+	}
+	filtered := make([]types.ImageInfo, 0, len(infos))
+	for _, info := range infos {
+		if urls[info.URL] || urls[info.OriginalURL] {
+			filtered = append(filtered, info)
+		}
+	}
+	return marshalImageInfos(filtered)
+}
+
+func marshalImageInfos(infos []types.ImageInfo) string {
+	if len(infos) == 0 {
+		return ""
+	}
+	data, err := json.Marshal(infos)
+	if err != nil {
+		return ""
+	}
+	return string(data)
+}
+
+// PruneMarkdownImagesOutsideRange removes Markdown image lines whose
+// document-level rune offsets fall outside [matchStart, matchEnd). Non-image
+// text is preserved so parent-child expansion still provides full textual
+// context while dropping irrelevant page thumbnails.
+func PruneMarkdownImagesOutsideRange(
+	content string,
+	contentStartAt, matchStart, matchEnd int,
+) string {
+	locs := MarkdownImageRegex.FindAllStringIndex(content, -1)
+	if len(locs) == 0 {
+		return content
+	}
+	var b strings.Builder
+	last := 0
+	for _, loc := range locs {
+		if loc[0] < last || loc[1] > len(content) {
+			continue
+		}
+		docStart := contentStartAt + utf8.RuneCountInString(content[:loc[0]])
+		docEnd := contentStartAt + utf8.RuneCountInString(content[:loc[1]])
+		inRange := docStart < matchEnd && docEnd > matchStart
+		if inRange {
+			b.WriteString(content[last:loc[1]])
+		} else {
+			b.WriteString(content[last:loc[0]])
+		}
+		last = loc[1]
+	}
+	b.WriteString(content[last:])
+	return collapseBlankLines(b.String())
+}
+
+func collapseBlankLines(s string) string {
+	for strings.Contains(s, "\n\n\n") {
+		s = strings.ReplaceAll(s, "\n\n\n", "\n\n")
+	}
+	return strings.TrimSpace(s)
+}
@@ -0,0 +1,95 @@
+package searchutil
+
+import (
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+func TestSliceContentByDocumentRange(t *testing.T) {
+	parent := "aaaPAGE1bbbPAGE2ccc"
+	got := SliceContentByDocumentRange(parent, 100, 103, 108)
+	want := "PAGE1"
+	if got != want {
+		t.Fatalf("slice: got %q, want %q", got, want)
+	}
+}
+
+func TestFilterImageInfoByMatchRange(t *testing.T) {
+	parent := "![p1](u1)\n\n![p2](u2)\n\n![p3](u3)"
+	matchStart := len([]rune("![p1](u1)\n\n"))
+	matchEnd := matchStart + len([]rune("![p2](u2)"))
+	all := []types.ImageInfo{
+		{URL: "u1"}, {URL: "u2"}, {URL: "u3"},
+	}
+	raw, err := json.Marshal(all)
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := FilterImageInfoByMatchRange(parent, 0, matchStart, matchEnd, string(raw))
+	var filtered []types.ImageInfo
+	if err := json.Unmarshal([]byte(got), &filtered); err != nil {
+		t.Fatal(err)
+	}
+	if len(filtered) != 1 || filtered[0].URL != "u2" {
+		t.Fatalf("filtered: %+v", filtered)
+	}
+}
+
+func TestFilterImageInfoByContentURLs(t *testing.T) {
+	content := "intro\n![page3](local://img3.jpg)\noutro"
+	all := []types.ImageInfo{
+		{URL: "local://img1.jpg", OCRText: "one"},
+		{URL: "local://img3.jpg", OCRText: "three"},
+	}
+	raw, err := json.Marshal(all)
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := FilterImageInfoByContentURLs(content, string(raw))
+	var filtered []types.ImageInfo
+	if err := json.Unmarshal([]byte(got), &filtered); err != nil {
+		t.Fatalf("unmarshal filtered: %v", err)
+	}
+	if len(filtered) != 1 || filtered[0].URL != "local://img3.jpg" {
+		t.Fatalf("filtered: %+v", filtered)
+	}
+}
+
+func TestPruneMarkdownImagesOutsideRange(t *testing.T) {
+	parent := "![p1](u1)\n\n![p2](u2)\n\n![p3](u3)"
+	matchStart := len([]rune("![p1](u1)\n\n"))
+	matchEnd := matchStart + len([]rune("![p2](u2)"))
+	got := PruneMarkdownImagesOutsideRange(parent, 0, matchStart, matchEnd)
+	if got != "![p2](u2)" {
+		t.Fatalf("prune: got %q", got)
+	}
+}
+
+func TestEnrichContentWithImageInfoForChat_SkipsUnmatched(t *testing.T) {
+	content := "![p1](u1)\n\n![p2](u2)"
+	raw, _ := json.Marshal([]types.ImageInfo{{URL: "u2", OCRText: "two"}})
+	got := EnrichContentWithImageInfoForChat(content, string(raw))
+	if strings.Count(got, "<image url=") != 1 {
+		t.Fatalf("expected one image block, got: %s", got)
+	}
+	if !strings.Contains(got, "![p1](u1)") {
+		t.Fatalf("unmatched markdown should remain: %s", got)
+	}
+	if !strings.Contains(got, "<image_ocr>two</image_ocr>") {
+		t.Fatalf("matched image should be enriched: %s", got)
+	}
+	if strings.Contains(got, "<image_original>") {
+		t.Fatalf("chat enrich should not duplicate markdown in image_original: %s", got)
+	}
+}
+
+func TestImageURLsInContent(t *testing.T) {
+	content := "![a](u1) x ![b](u2)"
+	urls := ImageURLsInContent(content)
+	if !urls["u1"] || !urls["u2"] || len(urls) != 2 {
+		t.Fatalf("urls: %#v", urls)
+	}
+}