diff --git a/internal/application/service/chat_pipeline/into_chat_message.go b/internal/application/service/chat_pipeline/into_chat_message.go index 9baec373..9df2cf70 100644 --- a/internal/application/service/chat_pipeline/into_chat_message.go +++ b/internal/application/service/chat_pipeline/into_chat_message.go @@ -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) } diff --git a/internal/application/service/chat_pipeline/merge.go b/internal/application/service/chat_pipeline/merge.go index ae199da5..82a6506d 100644 --- a/internal/application/service/chat_pipeline/merge.go +++ b/internal/application/service/chat_pipeline/merge.go @@ -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) + } } diff --git a/internal/application/service/chat_pipeline/merge_parent_resolve_test.go b/internal/application/service/chat_pipeline/merge_parent_resolve_test.go new file mode 100644 index 00000000..30f7ed57 --- /dev/null +++ b/internal/application/service/chat_pipeline/merge_parent_resolve_test.go @@ -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) + } +} diff --git a/internal/handler/session/helpers.go b/internal/handler/session/helpers.go index a0a1fda1..8cd1e5e9 100644 --- a/internal/handler/session/helpers.go +++ b/internal/handler/session/helpers.go @@ -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 { diff --git a/internal/searchutil/imageinfo.go b/internal/searchutil/imageinfo.go index ae9f9e97..5499781c 100644 --- a/internal/searchutil/imageinfo.go +++ b/internal/searchutil/imageinfo.go @@ -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 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("\n", imgURL)) + b.WriteString(inner) + b.WriteString("") + 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 diff --git a/internal/searchutil/imageinfo_match.go b/internal/searchutil/imageinfo_match.go new file mode 100644 index 00000000..d2c0280c --- /dev/null +++ b/internal/searchutil/imageinfo_match.go @@ -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) +} diff --git a/internal/searchutil/imageinfo_match_test.go b/internal/searchutil/imageinfo_match_test.go new file mode 100644 index 00000000..75fae7be --- /dev/null +++ b/internal/searchutil/imageinfo_match_test.go @@ -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, "two") { + t.Fatalf("matched image should be enriched: %s", got) + } + if strings.Contains(got, "") { + 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) + } +}