package knowledge_compiler import ( "context" "encoding/json" "errors" "fmt" "hash/fnv" "regexp" "strings" "sync" "testing" "ragflow/internal/ingestion/component/knowledge_compiler/common" ) // mockChat answers the structure variant's three LLM call shapes under the // Python-aligned items contract: node extraction (entity items), edge // extraction (relation items constrained to Known Entities), and merge // judging (Item A / Item B). It scans the prompt for the Greek-letter fixture // universe (Alpha..Epsilon) so assertions hold regardless of batch packing. type mockChat struct{} var greekNames = []string{"Alpha", "Beta", "Gamma", "Delta", "Epsilon"} var greekPairs = [][2]string{{"Alpha", "Beta"}, {"Beta", "Gamma"}, {"Gamma", "Delta"}, {"Delta", "Epsilon"}, {"Alpha", "Delta"}, {"Gamma", "Epsilon"}} // graphParserConfig declares the entity/relation template shape so the // structure variant compiles the hypergraph kind (mirrors a Python // parser_config whose kind aliases "graph" -> hypergraph). func graphParserConfig() map[string]any { return map[string]any{ "kind": "graph", "entity": map[string]any{"fields": []any{ map[string]any{"type": "project", "description": "a named project"}, }}, "relation": map[string]any{"fields": []any{ map[string]any{"type": "linked", "description": "a link between projects"}, }}, } } var chunkIDRe = regexp.MustCompile(`\[CHUNK_ID: ([^\]]+)\]`) func chunkIDsInPrompt(prompt string) []string { var ids []string for _, m := range chunkIDRe.FindAllStringSubmatch(prompt, -1) { ids = append(ids, m[1]) } if len(ids) == 0 { ids = []string{"c1"} } return ids } func itemsJSON(items []map[string]any, chunkIDs []string) string { for _, item := range items { item["source_chunk_ids"] = chunkIDs } b, _ := json.Marshal(map[string]any{"items": items}) return string(b) } func parseMergePayloads(prompt string) (map[string]any, map[string]any) { body := strings.TrimPrefix(prompt, "Item A (existing):\n") parts := strings.SplitN(body, "\n\nItem B (incoming):\n", 2) if len(parts) != 2 { return nil, nil } var a, b map[string]any if err := json.Unmarshal([]byte(parts[0]), &a); err != nil { return nil, nil } if err := json.Unmarshal([]byte(parts[1]), &b); err != nil { return nil, nil } return a, b } func (mockChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { if !req.JSONMode { // Non-JSON calls (mindmap outline, wiki prose paths) get plain text. return &common.ChatResponse{Content: "Outline of: " + firstWords(req.UserPrompt, 24)}, nil } switch { case strings.HasPrefix(req.UserPrompt, "Item A (existing):"): a, b := parseMergePayloads(req.UserPrompt) if a != nil && b != nil && a["name"] != nil && a["name"] == b["name"] { merged, _ := json.Marshal(a) return &common.ChatResponse{Content: fmt.Sprintf(`{"duplicated":true,"merged":%s}`, merged)}, nil } return &common.ChatResponse{Content: `{"duplicated":false,"merged":null}`}, nil case strings.Contains(req.SystemPrompt, "## Known Entities:"): var items []map[string]any for _, pr := range greekPairs { if strings.Contains(req.UserPrompt, pr[0]) && strings.Contains(req.UserPrompt, pr[1]) { items = append(items, map[string]any{"type": "linked", "source": pr[0], "target": pr[1], "description": pr[0] + " linked to " + pr[1]}) } } return &common.ChatResponse{Content: itemsJSON(items, chunkIDsInPrompt(req.UserPrompt))}, nil default: var items []map[string]any for _, n := range greekNames { if strings.Contains(req.UserPrompt, n) { items = append(items, map[string]any{"type": "project", "name": n, "description": n + " project"}) } } return &common.ChatResponse{Content: itemsJSON(items, chunkIDsInPrompt(req.UserPrompt))}, nil } } // proseChat returns generic, non-empty prose for the non-structure variants // (wiki/raptor/mindmap/datasetnav), which all just need a summary/outline text. type proseChat struct{} func (proseChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { if req.JSONMode { return &common.ChatResponse{Content: `{"ok":true}`}, nil } // Echo a deterministic prose reply derived from the prompt. return &common.ChatResponse{Content: "Summary of: " + firstWords(req.UserPrompt, 24)}, nil } func firstWords(s string, n int) string { fields := strings.Fields(s) if len(fields) > n { fields = fields[:n] } return strings.Join(fields, " ") } func installProseDeps(t *testing.T) { t.Helper() common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{Chat: proseChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID}, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) } type mockEmbedder struct{ dim int } func (m mockEmbedder) Dimensions() int { return m.dim } func (m mockEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { out := make([][]float32, len(texts)) for i, t := range texts { out[i] = deterministicVec(t, m.dim) } return out, nil } func deterministicVec(s string, dim int) []float32 { h := fnv.New32a() _, _ = h.Write([]byte(s)) seed := h.Sum32() v := make([]float32, dim) for j := 0; j < dim; j++ { v[j] = float32((seed>>uint(j*5))&0xFF) / 255.0 } return v } func installMockDeps(t *testing.T) { t.Helper() common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{Chat: mockChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID}, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) } func TestKnowledgeCompiler_Structure_EndToEnd(t *testing.T) { installMockDeps(t) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "Alpha is a Beta"}, map[string]any{"id": "c2", "text": "Beta related to Gamma"}, }, "doc_id": "d1", "tenant_id": "t1", "dataset_id": "ds1", "parser_config": graphParserConfig(), }) if err != nil { t.Fatalf("Invoke: %v", err) } chunks, ok := out["chunks"].([]any) if !ok { t.Fatalf("chunks = %T, want []any", out["chunks"]) } // 2 input chunks + 3 entities + 2 relations + 1 graph = 8 total (the // compiled knowledge units are merged into the upstream input chunks). if len(chunks) != 8 { t.Fatalf("len(chunks) = %d, want 8 (2 input + 3 entities + 2 relations + 1 graph)", len(chunks)) } // Exactly one graph product, parseable to {entities, relations} (mirrors // Python's _struct_rebuild_graph_json). var graph map[string]any for _, c := range chunks { cm, ok := c.(map[string]any) if !ok { continue } if kind, _ := cm["kc_kind"].(string); kind == "graph" { if err := json.Unmarshal([]byte(cm["text"].(string)), &graph); err != nil { t.Fatalf("unmarshal graph: %v", err) } } } if graph == nil { t.Fatal("no graph chunk emitted") } entities, _ := graph["entities"].([]any) relations, _ := graph["relations"].([]any) if len(entities) != 3 { t.Fatalf("graph entities = %d, want 3 (Alpha/Beta/Gamma)", len(entities)) } if len(relations) != 2 { t.Fatalf("graph relations = %d, want 2", len(relations)) } } func TestKnowledgeCompiler_UnknownVariant(t *testing.T) { c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{"variant": "nope"}) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } _, err = c.Invoke(context.Background(), nil, map[string]any{"chunks": []any{}}) if !errors.Is(err, common.ErrUnknownVariant) { t.Fatalf("err = %v, want ErrUnknownVariant", err) } } func TestKnowledgeCompiler_Alias_Mindmap(t *testing.T) { installMockDeps(t) // "mind_map" is the deprecated alias for "mindmap"; both resolve to the // implemented mindmap variant and must run (not ErrUnknownVariant / stub). c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "mind_map", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, "doc_id": "d1", "tenant_id": "t1", }) if err != nil { t.Fatalf("Invoke mind_map: %v", err) } chunks, ok := out["chunks"].([]any) if !ok || len(chunks) == 0 { t.Fatalf("mind_map produced no chunks: %v", out["chunks"]) } } // runVariant is a shared harness that builds a KnowledgeCompiler component for // the given variant and invokes it over a couple of chunks, returning the merged // chunk list (upstream input chunks + compiled knowledge-unit chunks). func runVariant(t *testing.T, variant string, extra map[string]any) []map[string]any { t.Helper() params := map[string]any{"variant": variant, "llm_id": "llm1", "embedding_model": "emb1"} for k, v := range extra { params[k] = v } c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", params) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent(%s): %v", variant, err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog near the river bank."}, map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river at dawn."}, map[string]any{"id": "c3", "text": "Rivers flow through valleys carrying water from the mountains."}, }, "doc_id": "d1", "tenant_id": "t1", }) if err != nil { t.Fatalf("Invoke(%s): %v", variant, err) } raw, ok := out["chunks"].([]any) if !ok { t.Fatalf("Invoke(%s): chunks = %T", variant, out["chunks"]) } if len(raw) == 0 { t.Fatalf("Invoke(%s): produced no chunks", variant) } // Every emitted chunk must carry an id; every compiled knowledge unit must // carry its embedding (q__vec) and a compile_kwd discriminator. chunks := make([]map[string]any, 0, len(raw)) seenCompiled := false for _, r := range raw { cm, ok := r.(map[string]any) if !ok { t.Fatalf("Invoke(%s): chunk not a map: %T", variant, r) } chunks = append(chunks, cm) if cm["id"] == nil || cm["id"] == "" { t.Fatalf("Invoke(%s): chunk missing id: %v", variant, cm) } if ck, _ := cm["compile_kwd"].(string); ck != "" { seenCompiled = true if _, ok := cm["q_8_vec"]; !ok { t.Fatalf("Invoke(%s): compiled chunk missing vector: %v", variant, cm) } } } if !seenCompiled { t.Fatalf("Invoke(%s): no compiled knowledge-unit chunks emitted", variant) } return chunks } func TestKnowledgeCompiler_Wiki_EndToEnd(t *testing.T) { installProseDeps(t) chunks := runVariant(t, "wiki", nil) // wiki produces a "page" chunk (kind stored as kc_kind). foundPage := false for _, c := range chunks { if kind, _ := c["kc_kind"].(string); kind == "page" { foundPage = true } } if !foundPage { t.Fatalf("wiki: no 'page' chunk; got %d chunks", len(chunks)) } } func TestKnowledgeCompiler_Raptor_EndToEnd(t *testing.T) { installProseDeps(t) // Psi builder (default): zero clustering dependency. chunks := runVariant(t, "raptor", map[string]any{"extra": map[string]any{"clustering_method": "PSI"}}) foundRoot := false for _, c := range chunks { if kind, _ := c["kc_kind"].(string); kind == "root" { foundRoot = true } } if !foundRoot { t.Fatalf("raptor(PSI): no 'root' chunk; got %d chunks", len(chunks)) } // AHC builder must also run (PCA + Ward). chunksAHC := runVariant(t, "raptor", map[string]any{"extra": map[string]any{"clustering_method": "AHC"}}) if len(chunksAHC) == 0 { t.Fatalf("raptor(AHC): produced no chunks") } // GMM backend (diagonal-GMM EM + BIC) is implemented and must run, // producing a well-formed tree (root present, chunks non-empty). chunksGMM := runVariant(t, "raptor", map[string]any{"extra": map[string]any{"clustering_method": "GMM"}}) if len(chunksGMM) == 0 { t.Fatalf("raptor(GMM): produced no chunks") } foundRootGMM := false for _, c := range chunksGMM { if kind, _ := c["kc_kind"].(string); kind == "root" { foundRootGMM = true } } if !foundRootGMM { t.Fatalf("raptor(GMM): no 'root' chunk; got %d chunks", len(chunksGMM)) } } func TestKnowledgeCompiler_Mindmap_EndToEnd(t *testing.T) { installProseDeps(t) // The proseChat reply is flat text; parseOutline still yields a root + the // reply as a child node, so chunks are non-empty and parent-linked. chunks := runVariant(t, "mindmap", nil) // Root chunk must have empty parent_kwd and kind "root". var root map[string]any for _, c := range chunks { if kind, _ := c["kc_kind"].(string); kind == "root" { root = c } } if root == nil { t.Fatalf("mindmap: no root chunk") } if pid, _ := root["parent_kwd"].(string); pid != "" { t.Fatalf("mindmap: root parent_kwd = %q, want empty", pid) } } func TestKnowledgeCompiler_Datasetnav_EndToEnd(t *testing.T) { installProseDeps(t) chunks := runVariant(t, "datasetnav", nil) foundRoot := false for _, c := range chunks { if kind, _ := c["kc_kind"].(string); kind == "root" { foundRoot = true } } if !foundRoot { t.Fatalf("datasetnav: no 'root' chunk; got %d chunks", len(chunks)) } } // TestKnowledgeCompiler_EmitsChunks verifies that after compiling, the // component returns the knowledge units merged into the chunk stream (no // separate products/writer seam). The output shape is the chunker's // {output_format:"chunks", chunks:[...]}. func TestKnowledgeCompiler_EmitsChunks(t *testing.T) { installMockDeps(t) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "Alpha is a Beta"}, map[string]any{"id": "c2", "text": "Beta related to Gamma"}, }, "doc_id": "d1", "tenant_id": "t1", "dataset_id": "ds1", "parser_config": graphParserConfig(), }) if err != nil { t.Fatalf("Invoke: %v", err) } if got, _ := out["output_format"].(string); got != "chunks" { t.Fatalf("output_format = %v, want chunks", out["output_format"]) } // No legacy products/written keys remain in the output surface. if _, ok := out["products"]; ok { t.Fatalf("legacy 'products' key still present in output") } if _, ok := out["written"]; ok { t.Fatalf("legacy 'written' key still present in output") } compiled := 0 for _, r := range out["chunks"].([]any) { cm := r.(map[string]any) // Structure rows carry the inferred compile kind (hypergraph here), // mirroring Python's per-row autotype stamp. if ck, _ := cm["compile_kwd"].(string); ck == "hypergraph" { compiled++ if cm["id"] == nil || cm["id"] == "" { t.Fatal("compiled chunk missing id") } if _, ok := cm["q_8_vec"]; !ok { t.Fatal("compiled chunk missing vector") } } } if compiled != 6 { t.Fatalf("compiled structure chunks = %d, want 6 (3 entities + 2 relations + 1 graph)", compiled) } } // TestKnowledgeCompiler_TemplateIDsAndProvenance verifies that // compilation_template_ids passed via the DSL params are stamped onto every // compiled chunk, and that structure entity rows carry source_chunk_ids // pointing back at their originating input chunks. func TestKnowledgeCompiler_TemplateIDsAndProvenance(t *testing.T) { installMockDeps(t) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", "compilation_template_ids": []any{"tpl-A", "tpl-B"}, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "Alpha is a Beta"}, map[string]any{"id": "c2", "text": "Beta related to Gamma"}, }, "doc_id": "d1", "tenant_id": "t1", "parser_config": graphParserConfig(), }) if err != nil { t.Fatalf("Invoke: %v", err) } for _, r := range out["chunks"].([]any) { cm := r.(map[string]any) if ck, _ := cm["compile_kwd"].(string); ck != "hypergraph" { continue } // Every compiled chunk must carry the resolved template ids. var tidsCount int switch tids := cm["compilation_template_ids"].(type) { case []any: tidsCount = len(tids) case []string: tidsCount = len(tids) } if tidsCount != 2 { t.Fatalf("compiled chunk %v: compilation_template_ids = %v, want 2", cm["id"], cm["compilation_template_ids"]) } // Entity rows must carry source_chunk_ids. if kg, _ := cm["knowledge_graph_kwd"].(string); kg == "entity" { var idsCount int switch ids := cm["source_chunk_ids"].(type) { case []any: idsCount = len(ids) case []string: idsCount = len(ids) } if idsCount == 0 { t.Fatalf("entity chunk %v: missing source_chunk_ids", cm["id"]) } } } } // constEmbedder returns one fixed vector for every input text. It is used to // make historical-dedup deterministic: every compiled product shares the same // vector, so a historical candidate carrying that vector is a guaranteed // near-duplicate. type constEmbedder struct { dim int vec []float32 } func (m constEmbedder) Dimensions() int { return m.dim } func (m constEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { out := make([][]float32, len(texts)) for i := range out { out[i] = m.vec } return out, nil } // sinkRecorder is a ChunkedSink that tallies streamed products, used to assert // the flush policy actually emits incrementally instead of holding the full set. type sinkRecorder struct { mu sync.Mutex total int batches int } func (s *sinkRecorder) Emit(_ context.Context, items []common.Product) error { s.mu.Lock() defer s.mu.Unlock() s.total += len(items) s.batches++ return nil } // TestKnowledgeCompiler_Raptor_DegenerateNoInfiniteLoop is a regression for the // High-1 bug: RAPTOR could recurse forever when a re-cluster returned a single // label covering all points. The default mockEmbedder emits non-negative // vectors, so under PSI (threshold=0) every pair has cosine >= 0 and the whole // corpus collapses into one cluster; with more than RecursionMinClusterSize // points the old code re-enqueued the identical work item at level+1 and hung. // This test uses 6 chunks and asserts the run terminates with a well-formed root. // (If the guard regresses, go test's timeout turns the hang into a failure.) func TestKnowledgeCompiler_Raptor_DegenerateNoInfiniteLoop(t *testing.T) { installProseDeps(t) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "raptor", "llm_id": "llm1", "embedding_model": "emb1", "extra": map[string]any{"clustering_method": "PSI"}, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "alpha beta gamma delta epsilon zeta"}, map[string]any{"id": "c2", "text": "one two three four five six"}, map[string]any{"id": "c3", "text": "red green blue yellow white black"}, map[string]any{"id": "c4", "text": "cat dog bird fish frog snake"}, map[string]any{"id": "c5", "text": "sun moon star planet comet meteor"}, map[string]any{"id": "c6", "text": "king queen prince duke earl count"}, }, "doc_id": "d1", "tenant_id": "t1", }) if err != nil { t.Fatalf("Invoke hung or errored (RAPTOR infinite recursion?): %v", err) } raw, ok := out["chunks"].([]any) if !ok { t.Fatalf("chunks = %T", out["chunks"]) } foundRoot := false for _, r := range raw { if cm, ok := r.(map[string]any); ok { if k, _ := cm["kc_kind"].(string); k == "root" { foundRoot = true } } } if !foundRoot { t.Fatalf("raptor(degenerate): no root chunk; got %d chunks", len(raw)) } } // TestKnowledgeCompiler_Wiki_HistoricalDedupDropsDuplicates is a regression for // the High-2 bug: EnableHistoricalDedup / the HistoricalCandidates override used // to be telemetry-only dead code. Here a constant embedder makes every wiki // product share one vector, and we supply that vector as a historical candidate; // the run must drop the near-duplicate products so none survive in the output. func TestKnowledgeCompiler_Wiki_HistoricalDedupDropsDuplicates(t *testing.T) { common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{ Chat: proseChat{}, Embed: constEmbedder{dim: 8, vec: []float32{1, 0, 0, 0, 0, 0, 0, 0}}, TenantID: tenantID, }, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "wiki", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog."}, map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river."}, }, "doc_id": "d1", "tenant_id": "t1", // Offline/test override: a historical candidate with the same vector as // every compiled product -> all are near-duplicates and must be dropped. "historical_candidates": []common.Candidate{ {ID: "old-artifact", Vector: []float32{1, 0, 0, 0, 0, 0, 0, 0}}, }, }) if err != nil { t.Fatalf("Invoke: %v", err) } raw, ok := out["chunks"].([]any) if !ok { t.Fatalf("chunks = %T", out["chunks"]) } for _, r := range raw { cm, ok := r.(map[string]any) if !ok { continue } if _, has := cm["compile_kwd"].(string); has { t.Fatalf("wiki historical dedup failed: a compiled chunk survived despite a matching historical candidate: %v", cm) } } } func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{ Chat: wikiUpdateChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID, WikiPages: wikiStoreTestStub{page: &common.WikiPageCandidate{ ID: "existing-1", Slug: "entity/alpha", Title: "Alpha", PageType: "entity", Topic: "Alpha", ContentMD: "# Alpha\n\nExisting fact.\n", ContentMDRaw: "# Alpha\n\nExisting fact.\n", EntityNames: []string{"Alpha"}, Score: 0.85, }}, }, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "wiki", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "Alpha launched a new process in 2026."}, }, "doc_id": "d1", "dataset_id": "ds1", "tenant_id": "t1", }) if err != nil { t.Fatalf("Invoke: %v", err) } raw, ok := out["chunks"].([]any) if !ok { t.Fatalf("chunks = %T", out["chunks"]) } var page map[string]any for _, r := range raw { cm, ok := r.(map[string]any) if !ok { continue } if cm["compile_kwd"] == "artifact_page" && cm["kc_kind"] == "page" && cm["slug_kwd"] == "entity/alpha" { page = cm break } } if page == nil { t.Fatalf("no merged wiki page found in output: %#v", raw) } text, _ := page["text"].(string) if !strings.Contains(text, "Existing fact.") || !strings.Contains(text, "Alpha launched a new process in 2026.") { t.Fatalf("merged page text = %q, want both existing and incoming facts", text) } if outlinks, ok := page["outlinks_kwd"].([]any); ok && len(outlinks) == 0 { t.Fatalf("outlinks_kwd should include rewritten see-also link: %#v", page) } } // TestKnowledgeCompiler_GuardrailFlushStreams is a regression for the Medium-1 // bug: the flush policy used to run only AFTER the full result set was // materialised, so it could not cap peak memory. With the ProductSink, products // are emitted to the sink as they are produced; here a tiny MaxItems forces // streaming, and we assert the sink received every product while the returned // chunk list holds only the upstream inputs (the compiled units went to the sink). func TestKnowledgeCompiler_GuardrailFlushStreams(t *testing.T) { installMockDeps(t) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } // Override guardrails to a tiny cap with the flush policy. (Guardrails are // not yet DSL-configurable; set directly on the exported Param.) kc, ok := c.(*KnowledgeCompilerComponent) if !ok { t.Fatalf("component is not *KnowledgeCompilerComponent: %T", c) } kc.Param.Guardrails = common.CapacityGuardrails{ MaxItems: 2, OnExceed: common.ExceedFlush, } sink := &sinkRecorder{} out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "Alpha is a Beta"}, map[string]any{"id": "c2", "text": "Beta related to Gamma"}, }, "doc_id": "d1", "tenant_id": "t1", "parser_config": graphParserConfig(), "sink": sink, }) if err != nil { t.Fatalf("Invoke: %v", err) } sink.mu.Lock() emitted, batches := sink.total, sink.batches sink.mu.Unlock() if batches == 0 || emitted == 0 { t.Fatalf("flush policy did not stream products to sink (batches=%d emitted=%d)", batches, emitted) } // The 6 compiled structure units were streamed to the sink; the returned // chunk list must therefore hold only the 2 upstream inputs. raw, ok := out["chunks"].([]any) if !ok { t.Fatalf("chunks = %T", out["chunks"]) } if len(raw) != 2 { t.Fatalf("after flush, returned chunks = %d, want 2 (upstream only; compiled streamed to sink)", len(raw)) } } // fakeHistoricalKNN records the datasetID it was queried with and optionally // returns a hit for every lookup, so a test can assert both the scope of the // KNN query and that near-duplicate products are dropped. type fakeHistoricalKNN struct { mu sync.Mutex lastDS string hit bool } func (f *fakeHistoricalKNN) TopKHistory(_ context.Context, _ string, datasetID, _ string, _ []float32, _ int, _ float64) ([]common.HistoricalHit, error) { f.mu.Lock() defer f.mu.Unlock() f.lastDS = datasetID if f.hit { return []common.HistoricalHit{{ID: "old", Score: 1.0}}, nil } return nil, nil } // TestKnowledgeCompiler_Wiki_HistoricalDedupScopedByDataset is a regression for // the Medium bug: the historical KNN lookup used to be scoped by doc_id, so // cross-document dedup within the same dataset did not work. This test supplies // both a doc_id ("d1") and a dataset_id ("ds1") and asserts the lookup is // scoped to the dataset, not the document. func TestKnowledgeCompiler_Wiki_HistoricalDedupScopedByDataset(t *testing.T) { knn := &fakeHistoricalKNN{hit: true} // every product is a near-dup -> dropped common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{ Chat: proseChat{}, Embed: constEmbedder{dim: 8, vec: []float32{1, 0, 0, 0, 0, 0, 0, 0}}, TenantID: tenantID, HistoricalKNN: knn, }, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "wiki", "llm_id": "llm1", "embedding_model": "emb1", "enable_historical_dedup": true, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog."}, map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river."}, }, "doc_id": "d1", "dataset_id": "ds1", "tenant_id": "t1", }) if err != nil { t.Fatalf("Invoke: %v", err) } knn.mu.Lock() ds := knn.lastDS knn.mu.Unlock() if ds != "ds1" { t.Fatalf("historical KNN queried with datasetID=%q, want %q (lookup must be scoped to the dataset, not the document)", ds, "ds1") } // With hit=true the near-duplicate products must be dropped, so no compiled // chunks survive. raw, ok := out["chunks"].([]any) if !ok { t.Fatalf("chunks = %T", out["chunks"]) } for _, r := range raw { if cm, ok := r.(map[string]any); ok { if _, has := cm["compile_kwd"]; has { t.Fatalf("wiki historical dedup failed: a compiled chunk survived despite a matching historical hit: %v", cm) } } } } // fakeDatasetnavLock records the key it was asked to acquire, so a test can // assert the datasetnav rebuild lock is scoped to the dataset, not the document. type fakeDatasetnavLock struct { mu sync.Mutex lastKey string } func (f *fakeDatasetnavLock) Acquire(_ context.Context, key string) (bool, error) { f.mu.Lock() defer f.mu.Unlock() f.lastKey = key return true, nil } func (f *fakeDatasetnavLock) Release(_ context.Context, _ string) error { return nil } // TestKnowledgeCompiler_Datasetnav_LockKeyedByDataset is a regression for the // Medium bug: the distributed rebuild lock used to be keyed by doc_id, so two // runs against the same dataset (with different per-document doc_id values) // took different locks and could interleave. This test supplies both doc_id // ("d1") and dataset_id ("ds1") and asserts the lock key uses the dataset. func TestKnowledgeCompiler_Datasetnav_LockKeyedByDataset(t *testing.T) { lk := &fakeDatasetnavLock{} common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{ Chat: proseChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID, Redis: lk, }, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "datasetnav", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } if _, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{ map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog near the river bank."}, map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river at dawn."}, }, "doc_id": "d1", "dataset_id": "ds1", "tenant_id": "t1", }); err != nil { t.Fatalf("Invoke: %v", err) } lk.mu.Lock() key := lk.lastKey lk.mu.Unlock() want := "datasetnav:t1:ds1" if key != want { t.Fatalf("datasetnav lock key = %q, want %q (lock must be scoped to the dataset, not the document)", key, want) } } // recordingNavChat echoes each nav-group summary back verbatim (so the child // summaries are identifiable) and records the root-synthesis user prompt so a // test can assert which child summaries reached the root overview. type recordingNavChat struct { mu sync.Mutex rootPrompt string } func (r *recordingNavChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { if strings.HasPrefix(req.UserPrompt, "Compose a navigation overview") { r.mu.Lock() r.rootPrompt = req.UserPrompt r.mu.Unlock() return &common.ChatResponse{Content: "root overview"}, nil } // Nav-group summary: echo the group text so the marker survives into the // root prompt when the root is built from all child summaries. return &common.ChatResponse{Content: strings.TrimSpace(req.UserPrompt)}, nil } // TestKnowledgeCompiler_Datasetnav_RootIncludesFlushedSummaries is a regression // for the Medium bug: the root overview used to be built from sink.Products() // (the not-yet-flushed buffer), so under OnExceed=flush any nav nodes already // emitted and cleared no longer contributed to the root, making the dataset // root incomplete precisely in the large-result case where streaming applies. // Here nav_radius > 1 forces each chunk into its own nav node, a MaxItems=1 // flush cap drains earlier nodes from the sink before the root is built, and we // assert every child marker still appears in the recorded root-synthesis prompt. func TestKnowledgeCompiler_Datasetnav_RootIncludesFlushedSummaries(t *testing.T) { chat := &recordingNavChat{} common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{ Chat: chat, Embed: mockEmbedder{dim: 8}, TenantID: tenantID, }, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "datasetnav", "llm_id": "llm1", "embedding_model": "emb1", // nav_radius > 1 guarantees no two chunks group together (cosine <= 1), // so each chunk becomes its own nav node. "extra": map[string]any{"nav_radius": 1.1}, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } kc, ok := c.(*KnowledgeCompilerComponent) if !ok { t.Fatalf("component is not *KnowledgeCompilerComponent: %T", c) } // Tiny cap + flush so earlier nav nodes are emitted and cleared from the // sink buffer before the root synthesis runs. kc.Param.Guardrails = common.CapacityGuardrails{MaxItems: 1, OnExceed: common.ExceedFlush} markers := []string{"NAVMARKA", "NAVMARKB", "NAVMARKC", "NAVMARKD", "NAVMARKE"} chunks := make([]any, len(markers)) for i, m := range markers { chunks[i] = map[string]any{"id": m, "text": m + " unique section body text"} } if _, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": chunks, "doc_id": "d1", "tenant_id": "t1", "sink": &sinkRecorder{}, }); err != nil { t.Fatalf("Invoke: %v", err) } chat.mu.Lock() root := chat.rootPrompt chat.mu.Unlock() if root == "" { t.Fatalf("root synthesis prompt was never recorded (root node not built)") } for _, m := range markers { if !strings.Contains(root, m) { t.Fatalf("root overview is missing flushed child summary %q; the root must be built from ALL child summaries, not just the not-yet-flushed sink buffer.\nroot prompt:\n%s", m, root) } } } // fencedChat wraps otherwise-valid extraction JSON in a ```json ... ``` fence, // the most common way LLMs wrap JSON even when JSONMode is requested. type fencedChat struct{} func (fencedChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { return &common.ChatResponse{Content: "```json\n" + `{"items":[{"type":"project","name":"Alpha","description":"Alpha project","source_chunk_ids":["c1"]}]}` + "\n```"}, nil } // proseOnlyChat returns genuine prose with no JSON at all — a model that // ignored the JSON contract entirely. type proseOnlyChat struct{} func (proseOnlyChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { return &common.ChatResponse{Content: "I could not extract structured facts from this passage."}, nil } type wikiUpdateChat struct{} func (wikiUpdateChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { switch { case req.JSONMode && strings.Contains(req.UserPrompt, "Extract all knowledge"): return &common.ChatResponse{Content: `{ "entities": [{"name":"Alpha","type":"person","aliases":["A"],"source_chunk_id":"c1"}], "concepts": [{"term":"Alpha Protocol","definition_excerpt":"Alpha Protocol is a rollout process","source_chunk_id":"c1"}], "claims": [{"statement":"Alpha launched a new process in 2026.","subject":"Alpha","confidence":"explicit","source_chunk_id":"c1"}], "relations": [], "topics": ["Alpha"] }`}, nil case req.JSONMode && strings.Contains(req.UserPrompt, "Return a JSON compilation plan"): return &common.ChatResponse{Content: `{ "pages":[{"action":"CREATE","slug":"entity/alpha","title":"Alpha","page_type":"entity","topic":"Alpha","entity_names":["Alpha"],"related_kb_pages":[],"priority":1,"lead":"Alpha overview","sections":[{"heading":"Overview","points":["Alpha overview"]}]}], "estimated_page_count":1, "compilation_notes":"ok" }`}, nil case strings.Contains(req.SystemPrompt, "wiki page merger"): return &common.ChatResponse{Content: "# Alpha\n\nExisting fact.\n\nAlpha launched a new process in 2026.\n\n## See also\n\n[Alpha Protocol](artifact/ds1/concept/alpha-protocol)"}, nil case strings.Contains(req.UserPrompt, "Existing page content"): return &common.ChatResponse{Content: "# Alpha\n\nAlpha launched a new process in 2026.\n\n## See also\n\n[[concept/alpha-protocol|Alpha Protocol]]"}, nil default: return &common.ChatResponse{Content: `{"action":"UPDATE","slug":"entity/alpha","reason":"same entity"}`}, nil } } type wikiStoreTestStub struct { page *common.WikiPageCandidate } func (s wikiStoreTestStub) FindSimilarPages(_ context.Context, _, _ string, _ []float32, _ int) ([]common.WikiPageCandidate, error) { if s.page == nil { return nil, nil } return []common.WikiPageCandidate{*s.page}, nil } func (s wikiStoreTestStub) GetPageBySlug(_ context.Context, _, _, slug string) (*common.WikiPageCandidate, error) { if s.page != nil && s.page.Slug == slug { return s.page, nil } return nil, nil } // TestKnowledgeCompiler_Structure_FencedJSONNotDropped is a regression for the // Medium bug: GenJSON used to return {"_raw":...} on any parse failure, which // parseExtractResult silently treated as "empty extraction", dropping every // entity/relation. A fenced ```json ... ``` reply is now unwrapped and parsed, // so the extraction still yields its entities. func TestKnowledgeCompiler_Structure_FencedJSONNotDropped(t *testing.T) { common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{Chat: fencedChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID}, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, "doc_id": "d1", "tenant_id": "t1", }) if err != nil { t.Fatalf("Invoke (fenced JSON should parse): %v", err) } var sawAlphaEntity bool for _, r := range out["chunks"].([]any) { cm := r.(map[string]any) if cm["compile_kwd"] == nil { continue } // name_kwd is lowercased (mirrors Python's _struct_to_doc_storage_doc). if cm["name_kwd"] == "alpha" { sawAlphaEntity = true } } if !sawAlphaEntity { t.Fatalf("fenced-JSON reply was dropped (no Alpha entity extracted); GenJSON must unwrap the fence, not treat it as empty extraction") } } // TestKnowledgeCompiler_Structure_MalformedJSONFailsLoud is the companion // regression: when the reply is genuinely unparseable (not just fenced), the // component must fail loudly instead of silently emitting zero knowledge units. func TestKnowledgeCompiler_Structure_MalformedJSONFailsLoud(t *testing.T) { common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{Chat: proseOnlyChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID}, nil }) t.Cleanup(func() { common.SetDepsResolver(nil) }) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } _, err = c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, "doc_id": "d1", "tenant_id": "t1", }) if err == nil { t.Fatalf("Invoke succeeded on an unparseable LLM reply; it must fail loudly rather than silently drop the extraction") } } // TestKnowledgeCompiler_PassThroughEnvelope is a regression for the Medium bug: // the advertised output schema promises name / tenant_id / kb_id are carried // forward from upstream (pass-through), but mergeChunks only emitted // output_format + chunks. Headless / manual chaining reads those identity keys // from the component output (the Tokenizer falls back to globals only in a full // pipeline), so they must be forwarded when present. func TestKnowledgeCompiler_PassThroughEnvelope(t *testing.T) { installMockDeps(t) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, "doc_id": "d1", "tenant_id": "t1", "name": "doc.pdf", "kb_id": "kb9", }) if err != nil { t.Fatalf("Invoke: %v", err) } if out["name"] != "doc.pdf" { t.Fatalf("output missing pass-through name: %v", out["name"]) } if out["tenant_id"] != "t1" { t.Fatalf("output missing pass-through tenant_id: %v", out["tenant_id"]) } if out["kb_id"] != "kb9" { t.Fatalf("output missing pass-through kb_id: %v", out["kb_id"]) } if out["output_format"] != "chunks" { t.Fatalf("output_format = %v, want chunks", out["output_format"]) } // Absent keys must NOT be invented. if _, ok := out["dataset_id"]; ok { t.Fatalf("output must not invent absent pass-through keys") } } // groupResolverStub maps every requested group id to a fixed pair of template // ids, standing in for the production DB-backed group service. func groupResolverStub(_ context.Context, _ string, groupIDs []string) ([]string, error) { var out []string for _, g := range groupIDs { out = append(out, "tpl-"+g+"-a", "tpl-"+g+"-b") } return out, nil } // TestKnowledgeCompiler_GroupIDsResolvedToTemplateIDs is a regression for the // open question: compilation_template_group_ids were parsed into Param.GroupIDs // but never resolved or stamped, so group-only configs silently missed // compilation_template_ids. The component now resolves groups via the installed // GroupResolver and merges them into the stamped template-id list. func TestKnowledgeCompiler_GroupIDsResolvedToTemplateIDs(t *testing.T) { installMockDeps(t) common.SetGroupResolver(groupResolverStub) t.Cleanup(func() { common.SetGroupResolver(nil) }) c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", "compilation_template_group_ids": []any{"grp1"}, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } out, err := c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, "doc_id": "d1", "tenant_id": "t1", }) if err != nil { t.Fatalf("Invoke: %v", err) } want := map[string]bool{"tpl-grp1-a": true, "tpl-grp1-b": true} checked := 0 for _, r := range out["chunks"].([]any) { cm := r.(map[string]any) // No parser_config is supplied, so InferType returns "list" and the // structure variant stamps compile_kwd="list" (not "structure"). if cm["compile_kwd"] != "list" { continue } checked++ var got []string switch v := cm["compilation_template_ids"].(type) { case []any: for _, e := range v { got = append(got, e.(string)) } case []string: got = v } if len(got) != 2 || !want[got[0]] || !want[got[1]] { t.Fatalf("compiled chunk %v: compilation_template_ids = %v, want %v (group ids must resolve to template ids)", cm["id"], got, want) } } if checked == 0 { t.Fatalf("no compiled chunks inspected (compile_kwd=list expected); assertion was vacuous") } } // TestKnowledgeCompiler_GroupIDsWithoutResolverFailsLoud is the companion // regression: a group-only config with no GroupResolver installed must fail // loudly (surfacing the misconfiguration) rather than silently emitting rows // that miss compilation_template_ids. func TestKnowledgeCompiler_GroupIDsWithoutResolverFailsLoud(t *testing.T) { installMockDeps(t) common.SetGroupResolver(nil) // ensure no resolver is installed c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", "compilation_template_group_ids": []any{"grp1"}, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } _, err = c.Invoke(context.Background(), nil, map[string]any{ "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, "doc_id": "d1", "tenant_id": "t1", }) if err == nil { t.Fatalf("Invoke succeeded with group ids but no GroupResolver; it must fail loudly instead of dropping compilation_template_ids") } if !strings.Contains(err.Error(), "GroupResolver") { t.Fatalf("error %q does not mention GroupResolver", err.Error()) } }