feat(vault): replace auto-linking with LLM-classified relationship types

Replace vector-similarity-only auto-linking (creating generic "semantic"
links) with an LLM classification step that determines actual relationship
types (reference, depends_on, extends, related, supersedes, contradicts)
and generates meaningful context descriptions.

- Add enrich_classify.go: orchestration, retry with escalating timeouts,
  candidate gathering with bidirectional dedup, 20-doc cap per batch
- Add enrich_classify_prompt.go: system prompt, JSON parsing with partial
  success model, UTF-8 safe truncation
- Restructure processBatch: summarize → embed → classify → dedup+wikilinks
- Move dedup recording after classify (failed classify allows re-enrichment)
- Remove autoLink method (fully replaced by classifyLinks)
- Add DeleteDocLinksByTypes to VaultStore interface + PG/SQLite (IN clause)
- Guard old link deletion behind len(newLinks) > 0 (no data loss on all-SKIP)
This commit is contained in:
viettranx committed 2026-04-10 09:54:02 +07:00
1 parent 6a0b4769d1
commit 4b7a2e6e58
12 files changed
+1819 -36

No files matched your search

+25
View File
@@ -4,6 +4,7 @@ import (
"context" "context"
"database/sql" "database/sql"
"fmt" "fmt"
"strings"
"time" "time"
"github.com/google/uuid" "github.com/google/uuid"
@@ -145,6 +146,30 @@ func (s *PGVaultStore) DeleteDocLinksByType(ctx context.Context, tenantID, docID
return err return err
} }
// DeleteDocLinksByTypes removes outbound links matching any of the given types from a document.
func (s *PGVaultStore) DeleteDocLinksByTypes(ctx context.Context, tenantID, docID string, types []string) error {
if len(types) == 0 {
return nil
}
uid := mustParseUUID(docID)
tid := mustParseUUID(tenantID)
// Build IN clause with positional params: $3, $4, ...
params := []any{uid, tid}
placeholders := make([]string, len(types))
for i, t := range types {
params = append(params, t)
placeholders[i] = fmt.Sprintf("$%d", i+3)
}
q := fmt.Sprintf(`
DELETE FROM vault_links vl
USING vault_documents vd
WHERE vl.from_doc_id = $1
AND vd.id = vl.from_doc_id AND vd.tenant_id = $2
AND vl.link_type IN (%s)`, strings.Join(placeholders, ","))
_, err := s.db.ExecContext(ctx, q, params...)
return err
}
func scanVaultLinks(rows *sql.Rows) ([]store.VaultLink, error) { func scanVaultLinks(rows *sql.Rows) ([]store.VaultLink, error) {
var links []store.VaultLink var links []store.VaultLink
for rows.Next() { for rows.Next() {
+24
View File
@@ -6,6 +6,7 @@ import (
"context" "context"
"database/sql" "database/sql"
"fmt" "fmt"
"strings"
"time" "time"
"github.com/google/uuid" "github.com/google/uuid"
@@ -122,6 +123,29 @@ func (s *SQLiteVaultStore) DeleteDocLinksByType(ctx context.Context, tenantID, d
return err return err
} }
// DeleteDocLinksByTypes removes outbound links matching any of the given types from a document.
func (s *SQLiteVaultStore) DeleteDocLinksByTypes(ctx context.Context, tenantID, docID string, types []string) error {
if len(types) == 0 {
return nil
}
// Build IN clause with ? placeholders.
params := []any{docID}
placeholders := make([]string, len(types))
for i, t := range types {
params = append(params, t)
placeholders[i] = "?"
}
params = append(params, tenantID)
q := fmt.Sprintf(`
DELETE FROM vault_links
WHERE from_doc_id = ?
AND link_type IN (%s)
AND from_doc_id IN (SELECT id FROM vault_documents WHERE tenant_id = ?)`,
strings.Join(placeholders, ","))
_, err := s.db.ExecContext(ctx, q, params...)
return err
}
func scanVaultLinkRows(rows *sql.Rows) ([]store.VaultLink, error) { func scanVaultLinkRows(rows *sql.Rows) ([]store.VaultLink, error) {
var links []store.VaultLink var links []store.VaultLink
for rows.Next() { for rows.Next() {
+1
View File
@@ -93,6 +93,7 @@ type VaultStore interface {
GetBacklinks(ctx context.Context, tenantID, docID string) ([]VaultBacklink, error) GetBacklinks(ctx context.Context, tenantID, docID string) ([]VaultBacklink, error)
DeleteDocLinks(ctx context.Context, tenantID, docID string) error DeleteDocLinks(ctx context.Context, tenantID, docID string) error
DeleteDocLinksByType(ctx context.Context, tenantID, docID, linkType string) error DeleteDocLinksByType(ctx context.Context, tenantID, docID, linkType string) error
DeleteDocLinksByTypes(ctx context.Context, tenantID, docID string, types []string) error
// Enrichment // Enrichment
// UpdateSummaryAndReembed updates summary text and re-generates embedding from title+path+summary. // UpdateSummaryAndReembed updates summary text and re-generates embedding from title+path+summary.
+203
View File
@@ -0,0 +1,203 @@
package vault
import (
"context"
"fmt"
"log/slog"
"maps"
"slices"
"strings"
"time"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
const (
classifyMaxRetries = 3
classifyMaxTokens = 1024
classifyTemperature = 0.1
classifyCtxMaxLen = 256 // max context string length stored in DB
classifySummaryMaxChars = 300 // max summary chars in prompt (validated: 300 for accuracy)
classifyMaxSourceDocs = 20 // max source docs per classifyLinks call (validated: cap unbounded time)
)
var classifyTimeouts = [3]time.Duration{5 * time.Minute, 7 * time.Minute, 10 * time.Minute}
var classifyBackoffs = [3]time.Duration{0, 2 * time.Second, 4 * time.Second}
// validClassifyTypes — accepted link types stored directly in DB (aligned with UI vault-link-dialog.tsx).
var validClassifyTypes = map[string]bool{
"reference": true, "depends_on": true, "extends": true,
"related": true, "supersedes": true, "contradicts": true,
}
type classifyDoc struct {
DocID, Title, Path, Summary string
}
type candidatePair struct {
Source, Candidate classifyDoc
Score float64
}
// classifyLinks orchestrates LLM-based link classification for enriched docs.
func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID string, results []enriched) {
if w.provider == nil {
return
}
capped := results
if len(capped) > classifyMaxSourceDocs {
capped = capped[:classifyMaxSourceDocs]
}
// On SQLite (desktop/Lite), FindSimilarDocs returns nil (no pgvector) — classify is a no-op.
candidates := w.gatherCandidates(ctx, tenantID, agentID, capped)
if len(candidates) == 0 {
return
}
allTypes := slices.Collect(maps.Keys(validClassifyTypes))
allTypes = append(allTypes, "semantic") // clean up legacy links
for sourceDocID, pairs := range candidates {
source := pairs[0].Source
candidateDocs := make([]classifyDoc, len(pairs))
for i, p := range pairs {
candidateDocs[i] = p.Candidate
}
system, user := buildClassifyPrompt(source, candidateDocs)
raw, err := w.callClassifyWithRetry(ctx, system, user)
if err != nil {
slog.Warn("vault.classify: llm_failed", "doc", sourceDocID, "err", err)
continue // SKIP fallback
}
parsed, err := parseClassifyResponse(raw, len(candidateDocs))
if err != nil {
hint := fmt.Sprintf("\n\nPrevious response was invalid JSON (error: %s). Output ONLY a valid JSON array.", err.Error())
raw2, err2 := w.callClassifyWithRetry(ctx, system, user+hint)
if err2 != nil {
slog.Warn("vault.classify: retry_parse_failed", "doc", sourceDocID, "err", err2)
continue
}
parsed, err = parseClassifyResponse(raw2, len(candidateDocs))
if err != nil {
slog.Warn("vault.classify: parse_still_failed", "doc", sourceDocID, "err", err)
continue
}
}
// Collect valid links (collect-then-write pattern).
var newLinks []store.VaultLink
for _, r := range parsed {
if r.Type == "SKIP" || !validClassifyTypes[r.Type] {
continue
}
linkCtx := r.Ctx
if len(linkCtx) > classifyCtxMaxLen {
linkCtx = string([]rune(linkCtx)[:classifyCtxMaxLen])
}
newLinks = append(newLinks, store.VaultLink{
FromDocID: sourceDocID,
ToDocID: candidateDocs[r.Idx-1].DocID, // idx is 1-based, validated by parseClassifyResponse
LinkType: r.Type,
Context: linkCtx,
})
}
// Only replace old links if LLM produced valid replacements (avoid data loss on all-SKIP).
if len(newLinks) > 0 {
if err := w.vault.DeleteDocLinksByTypes(ctx, tenantID, sourceDocID, allTypes); err != nil {
slog.Warn("vault.classify: delete_old", "doc", sourceDocID, "err", err)
}
for i := range newLinks {
if err := w.vault.CreateLink(ctx, &newLinks[i]); err != nil {
slog.Debug("vault.classify: create_link", "from", sourceDocID, "to", newLinks[i].ToDocID, "err", err)
}
}
}
}
}
func (w *enrichWorker) gatherCandidates(ctx context.Context, tenantID, agentID string, results []enriched) map[string][]candidatePair {
seen := make(map[string]bool)
out := make(map[string][]candidatePair)
for _, r := range results {
neighbors, err := w.vault.FindSimilarDocs(ctx, tenantID, agentID, r.payload.DocID, enrichSimilarityLimit)
if err != nil {
slog.Warn("vault.classify: find_similar", "doc", r.payload.DocID, "err", err)
continue
}
// Derive title from path (payload has no Title field; neighbor docs have it from DB).
title := r.payload.Path
if doc, err := w.vault.GetDocumentByID(ctx, tenantID, r.payload.DocID); err == nil && doc != nil {
title = doc.Title
}
src := classifyDoc{
DocID: r.payload.DocID,
Title: title,
Path: r.payload.Path,
Summary: truncateSummary(r.summary),
}
for _, n := range neighbors {
if n.Score < enrichSimilarityMin || n.Document.Summary == "" {
continue
}
// Bidirectional dedup: only process each pair once.
a, b := src.DocID, n.Document.ID
if a > b {
a, b = b, a
}
if key := a + ":" + b; seen[key] {
continue
} else {
seen[key] = true
}
out[src.DocID] = append(out[src.DocID], candidatePair{
Source: src,
Candidate: classifyDoc{
DocID: n.Document.ID,
Title: n.Document.Title,
Path: n.Document.Path,
Summary: truncateSummary(n.Document.Summary),
},
Score: n.Score,
})
}
}
return out
}
// callClassifyWithRetry calls the LLM with escalating timeouts and backoffs.
func (w *enrichWorker) callClassifyWithRetry(ctx context.Context, system, user string) (string, error) {
var lastErr error
for attempt := range classifyMaxRetries {
if attempt > 0 {
select {
case <-ctx.Done():
return "", ctx.Err()
case <-time.After(classifyBackoffs[attempt]):
}
}
cctx, cancel := context.WithTimeout(ctx, classifyTimeouts[attempt])
resp, err := w.provider.Chat(cctx, providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: system},
{Role: "user", Content: user},
},
Model: w.model,
Options: map[string]any{"max_tokens": classifyMaxTokens, "temperature": classifyTemperature},
})
cancel()
if err != nil {
lastErr = err
slog.Warn("vault.classify: retry", "attempt", attempt+1, "err", err)
continue
}
return strings.TrimSpace(resp.Content), nil
}
return "", fmt.Errorf("classify exhausted %d retries: %w", classifyMaxRetries, lastErr)
}
@@ -0,0 +1,186 @@
package vault
import (
"testing"
)
// ============================================================================
// JSON Parsing Edge Case Tests
// ============================================================================
// TestParseClassifyResponse_MixedValidInvalid handles mix of valid and invalid entries.
func TestParseClassifyResponse_MixedValidInvalid(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":"valid"},
{"idx":0,"type":"reference","ctx":"invalid idx"},
{"idx":2,"type":"bad_type","ctx":"invalid type"},
{"idx":3,"type":"extends","ctx":"valid"}
]`
results, err := parseClassifyResponse(raw, 3)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
// Should keep only idx 1 and 3 with valid types
if len(results) != 2 {
t.Fatalf("Expected 2 valid results, got %d", len(results))
}
if results[0].Idx != 1 || results[1].Idx != 3 {
t.Errorf("Invalid entries not filtered correctly")
}
}
// TestParseClassifyResponse_NegativeIdx filters silently.
func TestParseClassifyResponse_NegativeIdx(t *testing.T) {
raw := `[
{"idx":-1,"type":"reference","ctx":"invalid"},
{"idx":1,"type":"reference","ctx":"valid"}
]`
results, err := parseClassifyResponse(raw, 1)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 1 {
t.Fatalf("Expected 1 result (negative idx filtered), got %d", len(results))
}
if results[0].Idx != 1 {
t.Errorf("Negative idx should be filtered")
}
}
// TestParseClassifyResponse_AllSKIP preserves all SKIP entries.
func TestParseClassifyResponse_AllSKIP(t *testing.T) {
raw := `[
{"idx":1,"type":"SKIP","ctx":"no match"},
{"idx":2,"type":"SKIP","ctx":"different topic"}
]`
results, err := parseClassifyResponse(raw, 2)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 2 {
t.Fatalf("Expected 2 SKIP entries, got %d", len(results))
}
for i, r := range results {
if r.Type != "SKIP" {
t.Errorf("Result %d: expected SKIP, got %q", i, r.Type)
}
}
}
// TestParseClassifyResponse_LargeIdx beyond count are filtered.
func TestParseClassifyResponse_LargeIdx(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":"ok"},
{"idx":100,"type":"reference","ctx":"too large"},
{"idx":999,"type":"reference","ctx":"way too large"}
]`
results, err := parseClassifyResponse(raw, 5)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 1 {
t.Fatalf("Expected 1 result (oversized idx filtered), got %d", len(results))
}
if results[0].Idx != 1 {
t.Errorf("Large idx should be filtered")
}
}
// TestParseClassifyResponse_MaxBoundaryIdx at exactly count is valid.
func TestParseClassifyResponse_MaxBoundaryIdx(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":"first"},
{"idx":5,"type":"reference","ctx":"last valid"}
]`
results, err := parseClassifyResponse(raw, 5)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 2 {
t.Fatalf("Expected 2 results, got %d", len(results))
}
if results[1].Idx != 5 {
t.Errorf("Max boundary idx (5) should be valid when count=5")
}
}
// TestParseClassifyResponse_EmptyCtx preserves empty context string.
func TestParseClassifyResponse_EmptyCtx(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":""},
{"idx":2,"type":"related","ctx":"has context"}
]`
results, err := parseClassifyResponse(raw, 2)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 2 {
t.Fatalf("Expected 2 results, got %d", len(results))
}
if results[0].Ctx != "" {
t.Errorf("Empty ctx should be preserved, got %q", results[0].Ctx)
}
}
// TestParseClassifyResponse_WhitespaceCtx preserves whitespace-only context.
func TestParseClassifyResponse_WhitespaceCtx(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":" "},
{"idx":2,"type":"related","ctx":"valid"}
]`
results, err := parseClassifyResponse(raw, 2)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 2 {
t.Fatalf("Expected 2 results, got %d", len(results))
}
// Whitespace should be preserved
if results[0].Ctx != " " {
t.Errorf("Whitespace ctx should be preserved, got %q", results[0].Ctx)
}
}
// TestParseClassifyResponse_CaseInsensitiveType only SKIP and validClassifyTypes are accepted.
func TestParseClassifyResponse_CaseInsensitiveType(t *testing.T) {
// Lowercase versions of valid types should not be accepted
raw := `[
{"idx":1,"type":"reference","ctx":"ok"},
{"idx":2,"type":"REFERENCE","ctx":"uppercase"}
]`
results, err := parseClassifyResponse(raw, 2)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
// Only lowercase "reference" should be valid
if len(results) != 1 {
t.Fatalf("Expected 1 result (uppercase REFERENCE filtered), got %d", len(results))
}
if results[0].Type != "reference" {
t.Errorf("Only lowercase types accepted, got %q", results[0].Type)
}
}
@@ -0,0 +1,193 @@
package vault
import (
"fmt"
"strings"
"testing"
)
// ============================================================================
// JSON Parsing Tests
// ============================================================================
// TestParseClassifyResponse_Valid parses standard JSON array.
func TestParseClassifyResponse_Valid(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":"mentions config"},
{"idx":2,"type":"depends_on","ctx":"requires authentication"}
]`
results, err := parseClassifyResponse(raw, 2)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 2 {
t.Fatalf("Expected 2 results, got %d", len(results))
}
if results[0].Idx != 1 || results[0].Type != "reference" {
t.Errorf("Result 0: got idx=%d type=%q, want idx=1 type=reference", results[0].Idx, results[0].Type)
}
if results[1].Idx != 2 || results[1].Type != "depends_on" {
t.Errorf("Result 1: got idx=%d type=%q, want idx=2 type=depends_on", results[1].Idx, results[1].Type)
}
}
// TestParseClassifyResponse_WithCodeFence strips ```json fences.
func TestParseClassifyResponse_WithCodeFence(t *testing.T) {
raw := "```json\n[{\"idx\":1,\"type\":\"reference\",\"ctx\":\"test\"}]\n```"
results, err := parseClassifyResponse(raw, 1)
if err != nil {
t.Fatalf("parseClassifyResponse with code fence failed: %v", err)
}
if len(results) != 1 {
t.Fatalf("Expected 1 result, got %d", len(results))
}
if results[0].Type != "reference" {
t.Errorf("Expected type reference, got %q", results[0].Type)
}
}
// TestParseClassifyResponse_InvalidJSON returns error on unmarshal failure.
func TestParseClassifyResponse_InvalidJSON(t *testing.T) {
raw := `{invalid json}`
_, err := parseClassifyResponse(raw, 1)
if err == nil {
t.Fatalf("parseClassifyResponse should return error for invalid JSON")
}
if !strings.Contains(err.Error(), "unmarshal") {
t.Errorf("Error should mention unmarshal: %v", err)
}
}
// TestParseClassifyResponse_OutOfRangeIdx filters silently, keeps valid entries.
func TestParseClassifyResponse_OutOfRangeIdx(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":"ok"},
{"idx":99,"type":"related","ctx":"invalid idx"},
{"idx":2,"type":"extends","ctx":"ok"}
]`
results, err := parseClassifyResponse(raw, 2)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 2 {
t.Fatalf("Expected 2 results (out-of-range filtered), got %d", len(results))
}
if results[0].Idx != 1 || results[1].Idx != 2 {
t.Errorf("Out-of-range entry not filtered properly")
}
}
// TestParseClassifyResponse_UnknownType filters silently.
func TestParseClassifyResponse_UnknownType(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":"ok"},
{"idx":2,"type":"unknown_type","ctx":"invalid"},
{"idx":3,"type":"related","ctx":"ok"}
]`
results, err := parseClassifyResponse(raw, 3)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 2 {
t.Fatalf("Expected 2 results (unknown type filtered), got %d", len(results))
}
if results[0].Type != "reference" || results[1].Type != "related" {
t.Errorf("Unknown type not filtered properly")
}
}
// TestParseClassifyResponse_SKIPType preserves SKIP in output.
func TestParseClassifyResponse_SKIPType(t *testing.T) {
raw := `[
{"idx":1,"type":"reference","ctx":"valid"},
{"idx":2,"type":"SKIP","ctx":"no relationship"}
]`
results, err := parseClassifyResponse(raw, 2)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 2 {
t.Fatalf("Expected 2 results including SKIP, got %d", len(results))
}
skipFound := false
for _, r := range results {
if r.Type == "SKIP" {
skipFound = true
break
}
}
if !skipFound {
t.Errorf("SKIP type should be preserved in results")
}
}
// TestParseClassifyResponse_EmptyArray returns valid empty slice.
func TestParseClassifyResponse_EmptyArray(t *testing.T) {
raw := `[]`
results, err := parseClassifyResponse(raw, 5)
if err != nil {
t.Fatalf("parseClassifyResponse failed on empty array: %v", err)
}
if len(results) != 0 {
t.Fatalf("Expected 0 results, got %d", len(results))
}
}
// TestParseClassifyResponse_CtxTruncation truncates context over 256 chars.
func TestParseClassifyResponse_CtxTruncation(t *testing.T) {
longCtx := strings.Repeat("x", 300)
raw := fmt.Sprintf(`[{"idx":1,"type":"reference","ctx":"%s"}]`, longCtx)
results, err := parseClassifyResponse(raw, 1)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len([]rune(results[0].Ctx)) > classifyCtxMaxLen {
t.Errorf("Context not truncated: got %d runes, want ≤%d",
len([]rune(results[0].Ctx)), classifyCtxMaxLen)
}
}
// TestParseClassifyResponse_ZeroIdx filters silently (idx must be >= 1).
func TestParseClassifyResponse_ZeroIdx(t *testing.T) {
raw := `[
{"idx":0,"type":"reference","ctx":"invalid"},
{"idx":1,"type":"reference","ctx":"valid"}
]`
results, err := parseClassifyResponse(raw, 1)
if err != nil {
t.Fatalf("parseClassifyResponse failed: %v", err)
}
if len(results) != 1 {
t.Fatalf("Expected 1 result (idx 0 filtered), got %d", len(results))
}
if results[0].Idx != 1 {
t.Errorf("Zero idx should be filtered")
}
}
+90
View File
@@ -0,0 +1,90 @@
package vault
import (
"encoding/json"
"fmt"
"log/slog"
"strings"
)
// classifyResult is a single LLM classification output for a candidate doc.
type classifyResult struct {
Idx int `json:"idx"`
Type string `json:"type"`
Ctx string `json:"ctx"`
}
const classifySystemPrompt = `You classify relationships between documents in a knowledge vault.
## Link Types
- reference: A cites, mentions, or quotes B
- depends_on: A requires B to function (config, data dependency, prerequisite)
- extends: A expands, details, or elaborates on B
- related: A and B share the same topic or domain (use only when no stronger type fits)
- supersedes: A is a newer version that replaces or updates B
- contradicts: A conflicts with or opposes B's content
## Rules
- Output ONLY a valid JSON array, no other text
- Use SKIP when documents are similar but have no meaningful relationship
- Prefer specific types (reference, depends_on) over generic (related)
- Each candidate classified independently
- Keep ctx descriptions under 60 words
## Output Format
[{"idx":1,"type":"reference","ctx":"mentions OAuth config for setup"},{"idx":2,"type":"SKIP"}]`
// buildClassifyPrompt formats the system and user prompts for classify LLM call.
func buildClassifyPrompt(source classifyDoc, candidates []classifyDoc) (system, user string) {
var b strings.Builder
b.WriteString("## Source Document\n")
fmt.Fprintf(&b, "Title: %s\nPath: %s\nSummary: %s\n\n", source.Title, source.Path, source.Summary)
b.WriteString("## Candidates\n")
for i, c := range candidates {
fmt.Fprintf(&b, "%d. Title: %s | Path: %s | Summary: %s\n", i+1, c.Title, c.Path, c.Summary)
}
return classifySystemPrompt, b.String()
}
// parseClassifyResponse parses LLM JSON output into classify results.
// Uses partial success model: invalid entries filtered silently, error only on total unmarshal failure.
func parseClassifyResponse(raw string, count int) ([]classifyResult, error) {
raw = strings.TrimSpace(raw)
// Strip code fences.
raw = strings.TrimPrefix(raw, "```json")
raw = strings.TrimPrefix(raw, "```")
raw = strings.TrimSuffix(raw, "```")
raw = strings.TrimSpace(raw)
var results []classifyResult
if err := json.Unmarshal([]byte(raw), &results); err != nil {
return nil, fmt.Errorf("json unmarshal: %w", err)
}
// Filter invalid entries (partial success).
valid := results[:0]
for _, r := range results {
if r.Idx < 1 || r.Idx > count {
slog.Debug("vault.classify: idx out of range", "idx", r.Idx, "count", count)
continue
}
if r.Type != "SKIP" && !validClassifyTypes[r.Type] {
continue
}
// Truncate context.
if len(r.Ctx) > classifyCtxMaxLen {
r.Ctx = string([]rune(r.Ctx)[:classifyCtxMaxLen])
}
valid = append(valid, r)
}
return valid, nil
}
// truncateSummary caps summary at classifySummaryMaxChars with UTF-8 safe rune truncation.
func truncateSummary(s string) string {
runes := []rune(s)
if len(runes) <= classifySummaryMaxChars {
return s
}
return string(runes[:classifySummaryMaxChars]) + "..."
}
@@ -0,0 +1,148 @@
package vault
import (
"strings"
"testing"
)
// ============================================================================
// Prompt Building Tests
// ============================================================================
// TestBuildClassifyPrompt verifies system prompt contains all 6 types and user prompt has numbered candidates.
func TestBuildClassifyPrompt(t *testing.T) {
source := classifyDoc{
DocID: "doc1",
Title: "Test Document",
Path: "docs/test.md",
Summary: "This is a test document.",
}
candidates := []classifyDoc{
{
DocID: "doc2",
Title: "Related Doc 1",
Path: "docs/related1.md",
Summary: "First related document.",
},
{
DocID: "doc3",
Title: "Related Doc 2",
Path: "docs/related2.md",
Summary: "Second related document.",
},
}
system, user := buildClassifyPrompt(source, candidates)
// System prompt must contain all 6 relationship types
expectedTypes := []string{"reference", "depends_on", "extends", "related", "supersedes", "contradicts"}
for _, typ := range expectedTypes {
if !strings.Contains(system, typ) {
t.Errorf("System prompt missing type: %s", typ)
}
}
// System prompt must contain classifySystemPrompt constant text
if !strings.Contains(system, "Link Types") {
t.Errorf("System prompt missing 'Link Types' header")
}
// User prompt must contain source doc info
if !strings.Contains(user, "Test Document") {
t.Errorf("User prompt missing source title")
}
if !strings.Contains(user, "docs/test.md") {
t.Errorf("User prompt missing source path")
}
if !strings.Contains(user, "This is a test document") {
t.Errorf("User prompt missing source summary")
}
// User prompt must contain candidates with numbers (1., 2.)
if !strings.Contains(user, "1.") {
t.Errorf("User prompt missing numbered candidate (1.)")
}
if !strings.Contains(user, "2.") {
t.Errorf("User prompt missing numbered candidate (2.)")
}
// User prompt must contain candidate info
if !strings.Contains(user, "Related Doc 1") {
t.Errorf("User prompt missing candidate 1 title")
}
if !strings.Contains(user, "Related Doc 2") {
t.Errorf("User prompt missing candidate 2 title")
}
}
// ============================================================================
// Truncate Summary Tests
// ============================================================================
// TestTruncateSummary_Long truncates at classifySummaryMaxChars (300) with "..." suffix.
func TestTruncateSummary_Long(t *testing.T) {
// Create string longer than 300 chars
longText := strings.Repeat("a", 350)
result := truncateSummary(longText)
if len([]rune(result)) > classifySummaryMaxChars+3 { // +3 for "..."
t.Errorf("truncateSummary(%d chars) returned %d runes, want ≤%d",
len([]rune(longText)), len([]rune(result)), classifySummaryMaxChars+3)
}
if !strings.HasSuffix(result, "...") {
t.Errorf("Truncated result must end with '...', got: %q", result)
}
// Should be exactly 303 runes (300 + "...")
expected := len([]rune(strings.Repeat("a", classifySummaryMaxChars))) + 3
if len([]rune(result)) != expected {
t.Errorf("Truncated result has %d runes, want %d", len([]rune(result)), expected)
}
}
// TestTruncateSummary_Short returns short string as-is.
func TestTruncateSummary_Short(t *testing.T) {
short := "This is short"
result := truncateSummary(short)
if result != short {
t.Errorf("truncateSummary(%q) = %q, want unchanged", short, result)
}
if strings.Contains(result, "...") {
t.Errorf("Short string should not have '...' suffix: %q", result)
}
}
// TestTruncateSummary_BoundaryAt300 tests string exactly at 300 chars.
func TestTruncateSummary_BoundaryAt300(t *testing.T) {
exact := strings.Repeat("x", classifySummaryMaxChars)
result := truncateSummary(exact)
if result != exact {
t.Errorf("String at exactly %d chars should not be truncated", classifySummaryMaxChars)
}
if strings.HasSuffix(result, "...") {
t.Errorf("String at boundary should not have '...' suffix")
}
}
// TestTruncateSummary_Unicode tests UTF-8 safe truncation with multi-byte chars.
func TestTruncateSummary_Unicode(t *testing.T) {
// Use emoji (4 bytes each) to test multi-byte handling
longUnicode := strings.Repeat("🎉", 400) // 400 emoji chars, very long when encoded
result := truncateSummary(longUnicode)
runes := []rune(result)
if len(runes) > classifySummaryMaxChars+3 {
t.Errorf("Unicode truncation exceeded max: got %d runes, want ≤%d",
len(runes), classifySummaryMaxChars+3)
}
if !strings.HasSuffix(result, "...") {
t.Errorf("Truncated unicode must end with '...'")
}
}
@@ -0,0 +1,154 @@
package vault
import (
"context"
"errors"
"strings"
"testing"
)
// ============================================================================
// Advanced Retry Logic Tests
// ============================================================================
// TestCallClassifyWithRetry_FirstAttemptSuccess uses first timeout.
func TestCallClassifyWithRetry_FirstAttemptSuccess(t *testing.T) {
// Verify that first attempt uses classifyTimeouts[0]
if classifyTimeouts[0] == 0 {
t.Errorf("First timeout should be non-zero")
}
provider := &mockClassifyProvider{
responses: []string{`[{"idx":1,"type":"reference","ctx":"first attempt"}]`},
errors: []error{nil},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
if err != nil {
t.Fatalf("First attempt should succeed: %v", err)
}
if provider.calls != 1 {
t.Errorf("Expected exactly 1 call on first-attempt success, got %d", provider.calls)
}
if !strings.Contains(resp, "reference") {
t.Errorf("Response missing expected content: %q", resp)
}
}
// TestCallClassifyWithRetry_ResponseWhitespaceStripping trims whitespace from response.
func TestCallClassifyWithRetry_ResponseWhitespaceStripping(t *testing.T) {
// Response has leading/trailing whitespace that should be stripped
provider := &mockClassifyProvider{
responses: []string{"\n\n [{\"idx\":1,\"type\":\"reference\",\"ctx\":\"test\"}] \n"},
errors: []error{nil},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
if err != nil {
t.Fatalf("callClassifyWithRetry failed: %v", err)
}
// Response should be trimmed
if strings.HasPrefix(resp, "\n") || strings.HasSuffix(resp, "\n") {
t.Errorf("Response should be trimmed: %q", resp)
}
if !strings.Contains(resp, "reference") {
t.Errorf("Response missing expected content: %q", resp)
}
}
// TestCallClassifyWithRetry_MaxRetriesConstant verifies retry limit.
func TestCallClassifyWithRetry_MaxRetriesConstant(t *testing.T) {
if classifyMaxRetries != 3 {
t.Errorf("classifyMaxRetries should be 3, got %d", classifyMaxRetries)
}
// Verify arrays have correct length
if len(classifyTimeouts) != classifyMaxRetries {
t.Errorf("classifyTimeouts length should be %d, got %d", classifyMaxRetries, len(classifyTimeouts))
}
if len(classifyBackoffs) != classifyMaxRetries {
t.Errorf("classifyBackoffs length should be %d, got %d", classifyMaxRetries, len(classifyBackoffs))
}
}
// TestCallClassifyWithRetry_SecondAttemptSucceeds verifies first retry succeeds.
func TestCallClassifyWithRetry_SecondAttemptSucceeds(t *testing.T) {
provider := &mockClassifyProvider{
responses: []string{
"", // attempt 0: error
`[{"idx":1,"type":"extends","ctx":"second attempt"}]`, // attempt 1: success
},
errors: []error{
errors.New("first attempt failed"),
nil, // no error on second attempt
},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
if err != nil {
t.Fatalf("Should succeed on second attempt, got error: %v", err)
}
if provider.calls != 2 {
t.Errorf("Expected 2 calls, got %d", provider.calls)
}
if !strings.Contains(resp, "extends") {
t.Errorf("Response from second attempt not found: %q", resp)
}
}
// TestCallClassifyWithRetry_EmptyResponse returns error for empty response on all attempts.
func TestCallClassifyWithRetry_EmptyResponse(t *testing.T) {
provider := &mockClassifyProvider{
responses: []string{"", "", ""},
errors: []error{nil, nil, nil},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
// Empty response is still a successful LLM call, should return empty string
if err != nil {
t.Fatalf("Empty response should not error: %v", err)
}
if resp != "" {
t.Errorf("Expected empty response, got %q", resp)
}
if provider.calls != 1 {
t.Errorf("Expected 1 call (succeeded immediately), got %d", provider.calls)
}
}
@@ -0,0 +1,195 @@
package vault
import (
"context"
"fmt"
"strings"
"testing"
"github.com/nextlevelbuilder/goclaw/internal/providers"
)
// ============================================================================
// Mock Provider for Testing
// ============================================================================
// mockClassifyProvider implements providers.Provider for testing retry logic.
type mockClassifyProvider struct {
responses []string
errors []error
calls int
}
func (m *mockClassifyProvider) Chat(ctx context.Context, req providers.ChatRequest) (*providers.ChatResponse, error) {
idx := m.calls
m.calls++
if idx < len(m.errors) && m.errors[idx] != nil {
return nil, m.errors[idx]
}
resp := ""
if idx < len(m.responses) {
resp = m.responses[idx]
}
return &providers.ChatResponse{Content: resp}, nil
}
func (m *mockClassifyProvider) ChatStream(ctx context.Context, req providers.ChatRequest, onChunk func(providers.StreamChunk)) (*providers.ChatResponse, error) {
return m.Chat(ctx, req)
}
func (m *mockClassifyProvider) DefaultModel() string { return "test" }
func (m *mockClassifyProvider) Name() string { return "mock" }
// ============================================================================
// Retry Logic Tests
// ============================================================================
// TestCallClassifyWithRetry_Success succeeds on first attempt.
func TestCallClassifyWithRetry_Success(t *testing.T) {
provider := &mockClassifyProvider{
responses: []string{`[{"idx":1,"type":"reference","ctx":"test"}]`},
errors: []error{nil},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
if err != nil {
t.Fatalf("callClassifyWithRetry failed: %v", err)
}
if !strings.Contains(resp, "reference") {
t.Errorf("Response doesn't contain expected content: %q", resp)
}
if provider.calls != 1 {
t.Errorf("Expected 1 call, got %d", provider.calls)
}
}
// TestCallClassifyWithRetry_RetryThenSuccess fails twice, succeeds on third attempt.
func TestCallClassifyWithRetry_RetryThenSuccess(t *testing.T) {
provider := &mockClassifyProvider{
responses: []string{
"", // attempt 0: error
"", // attempt 1: error
`[{"idx":1,"type":"reference","ctx":"success"}]`, // attempt 2: success
},
errors: []error{
fmt.Errorf("network error"),
fmt.Errorf("timeout"),
nil, // no error on third attempt
},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
if err != nil {
t.Fatalf("callClassifyWithRetry should succeed after retries, got error: %v", err)
}
if !strings.Contains(resp, "success") {
t.Errorf("Response doesn't contain expected content: %q", resp)
}
if provider.calls != 3 {
t.Errorf("Expected 3 calls, got %d", provider.calls)
}
}
// TestCallClassifyWithRetry_AllFail exhausts retries and returns error.
func TestCallClassifyWithRetry_AllFail(t *testing.T) {
provider := &mockClassifyProvider{
responses: []string{"", "", ""},
errors: []error{
fmt.Errorf("error1"),
fmt.Errorf("error2"),
fmt.Errorf("error3"),
},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
ctx := context.Background()
_, err := worker.callClassifyWithRetry(ctx, "system", "user")
if err == nil {
t.Fatalf("callClassifyWithRetry should return error after exhausting retries")
}
if !strings.Contains(err.Error(), "exhausted") {
t.Errorf("Error should mention exhausted retries: %v", err)
}
if provider.calls != classifyMaxRetries {
t.Errorf("Expected %d calls, got %d", classifyMaxRetries, provider.calls)
}
}
// TestCallClassifyWithRetry_ContextCancellation respects context cancellation.
func TestCallClassifyWithRetry_ContextCancellation(t *testing.T) {
provider := &mockClassifyProvider{
responses: []string{"", "", ""},
errors: []error{
fmt.Errorf("error1"),
fmt.Errorf("error2"),
fmt.Errorf("error3"),
},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
ctx, cancel := context.WithCancel(context.Background())
cancel() // Cancel immediately
_, err := worker.callClassifyWithRetry(ctx, "system", "user")
if err == nil {
t.Fatalf("callClassifyWithRetry should return error for cancelled context")
}
if !strings.Contains(err.Error(), "context") {
t.Errorf("Error should mention context: %v", err)
}
}
// TestCallClassifyWithRetry_RetriesAndBackoffs verifies retry timeouts and backoffs escalate.
func TestCallClassifyWithRetry_RetriesAndBackoffs(t *testing.T) {
// Verify first backoff is 0 (no backoff on first attempt)
if classifyBackoffs[0] != 0 {
t.Errorf("First backoff should be 0, got %v", classifyBackoffs[0])
}
// Verify backoffs escalate for retry attempts
if classifyBackoffs[1] == 0 || classifyBackoffs[2] == 0 {
t.Errorf("Backoffs should escalate: %v", classifyBackoffs)
}
if classifyBackoffs[1] >= classifyBackoffs[2] {
t.Errorf("Backoffs should increase: %v", classifyBackoffs)
}
// Verify timeouts escalate
if classifyTimeouts[0] >= classifyTimeouts[1] || classifyTimeouts[1] >= classifyTimeouts[2] {
t.Errorf("Timeouts should escalate: %v", classifyTimeouts)
}
}
+22 -36
View File
@@ -2,7 +2,6 @@ package vault
import ( import (
"context" "context"
"fmt"
"log/slog" "log/slog"
"os" "os"
"path/filepath" "path/filepath"
@@ -80,6 +79,12 @@ func (w *enrichWorker) Handle(ctx context.Context, event eventbus.DomainEvent) e
return nil return nil
} }
// enriched holds a successfully summarized vault document pending embed+link.
type enriched struct {
payload eventbus.VaultDocUpsertedPayload
summary string
}
// processBatch drains and processes queued vault doc events in a loop. // processBatch drains and processes queued vault doc events in a loop.
func (w *enrichWorker) processBatch(ctx context.Context, key string) { func (w *enrichWorker) processBatch(ctx context.Context, key string) {
for { for {
@@ -91,10 +96,6 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) {
continue continue
} }
type enriched struct {
payload eventbus.VaultDocUpsertedPayload
summary string
}
var results []enriched var results []enriched
for _, item := range items { for _, item := range items {
@@ -145,20 +146,28 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) {
results = append(results, enriched{payload: item, summary: summary}) results = append(results, enriched{payload: item, summary: summary})
} }
// Update summary + embed + auto-link for each enriched doc. // Phase 2 — Embed: update summary + embed for all results.
// Do NOT record dedup here (moved to Phase 4 after classify).
var embedded []enriched
for _, r := range results { for _, r := range results {
if err := w.vault.UpdateSummaryAndReembed(ctx, r.payload.TenantID, r.payload.DocID, r.summary); err != nil { if err := w.vault.UpdateSummaryAndReembed(ctx, r.payload.TenantID, r.payload.DocID, r.summary); err != nil {
slog.Warn("vault.enrich: update_summary", "doc", r.payload.DocID, "err", err) slog.Warn("vault.enrich: update_summary", "doc", r.payload.DocID, "err", err)
continue // don't record in dedup continue
} }
embedded = append(embedded, r)
}
// Record hash only after successful update. // Phase 3 — Classify links (replaces autoLink).
// All items share same tenantID:agentID (batch queue key guarantees this).
if len(embedded) > 0 {
first := embedded[0].payload
w.classifyLinks(ctx, first.TenantID, first.AgentID, embedded)
}
// Phase 4 — Record dedup + wikilinks.
// Dedup recorded AFTER classify so failed classify allows re-enrichment.
for _, r := range embedded {
w.recordDedup(r.payload.DocID, r.payload.ContentHash) w.recordDedup(r.payload.DocID, r.payload.ContentHash)
// Auto-link via vector similarity.
w.autoLink(ctx, r.payload.TenantID, r.payload.AgentID, r.payload.DocID)
// Sync [[wikilinks]] from document content (text files only).
w.syncWikilinks(ctx, r.payload) w.syncWikilinks(ctx, r.payload)
} }
@@ -210,29 +219,6 @@ func (w *enrichWorker) syncWikilinks(ctx context.Context, p eventbus.VaultDocUps
} }
} }
// autoLink finds similar documents and creates semantic links.
func (w *enrichWorker) autoLink(ctx context.Context, tenantID, agentID, docID string) {
neighbors, err := w.vault.FindSimilarDocs(ctx, tenantID, agentID, docID, enrichSimilarityLimit)
if err != nil {
slog.Warn("vault.enrich: find_similar", "doc", docID, "err", err)
return
}
for _, n := range neighbors {
if n.Score < enrichSimilarityMin {
continue
}
link := &store.VaultLink{
FromDocID: docID,
ToDocID: n.Document.ID,
LinkType: "semantic",
Context: fmt.Sprintf("auto-linked (score: %.2f)", n.Score),
}
if err := w.vault.CreateLink(ctx, link); err != nil {
slog.Debug("vault.enrich: create_link", "from", docID, "to", n.Document.ID, "err", err)
}
}
}
// recordDedup stores a processed hash and evicts ~25% entries if over capacity. // recordDedup stores a processed hash and evicts ~25% entries if over capacity.
func (w *enrichWorker) recordDedup(docID, hash string) { func (w *enrichWorker) recordDedup(docID, hash string) {
+578
View File
@@ -0,0 +1,578 @@
//go:build integration
package integration
import (
"database/sql"
"errors"
"testing"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/store/pg"
)
func TestVaultClassify_DeleteDocLinksByTypes_SingleType(t *testing.T) {
db := testDB(t)
tenantID, agentID := seedTenantAgent(t, db)
ctx := tenantCtx(tenantID)
vs := newVaultStore(db)
tid := tenantID.String()
aid := agentID.String()
// Create 2 documents
docA := makeVaultDoc(tid, aid, "classify/source.md", "Source Doc")
docB := makeVaultDoc(tid, aid, "classify/target.md", "Target Doc")
if err := vs.UpsertDocument(ctx, docA); err != nil {
t.Fatalf("UpsertDocument A: %v", err)
}
if err := vs.UpsertDocument(ctx, docB); err != nil {
t.Fatalf("UpsertDocument B: %v", err)
}
// Create 3 links with different types
links := []struct {
linkType string
context string
}{
{"reference", "cited in source"},
{"wikilink", "mentioned in source"},
{"depends_on", "dependency of source"},
}
for _, l := range links {
link := &store.VaultLink{
FromDocID: docA.ID,
ToDocID: docB.ID,
LinkType: l.linkType,
Context: l.context,
}
if err := vs.CreateLink(ctx, link); err != nil {
t.Fatalf("CreateLink %s: %v", l.linkType, err)
}
}
// Verify: 3 links exist
outLinks, err := vs.GetOutLinks(ctx, tid, docA.ID)
if err != nil {
t.Fatalf("GetOutLinks before delete: %v", err)
}
if len(outLinks) != 3 {
t.Errorf("expected 3 links before delete, got %d", len(outLinks))
}
// Delete links with type "reference" only
if err := vs.DeleteDocLinksByTypes(ctx, tid, docA.ID, []string{"reference"}); err != nil {
t.Fatalf("DeleteDocLinksByTypes: %v", err)
}
// Verify: only "reference" deleted, "wikilink" and "depends_on" remain
outLinksAfter, err := vs.GetOutLinks(ctx, tid, docA.ID)
if err != nil {
t.Fatalf("GetOutLinks after delete: %v", err)
}
if len(outLinksAfter) != 2 {
t.Errorf("expected 2 links after deleting 1, got %d", len(outLinksAfter))
}
// Verify remaining links are correct types
typeMap := make(map[string]bool)
for _, l := range outLinksAfter {
typeMap[l.LinkType] = true
}
if !typeMap["wikilink"] || !typeMap["depends_on"] {
t.Errorf("expected wikilink and depends_on, got types: %v", typeMap)
}
if typeMap["reference"] {
t.Errorf("reference should have been deleted")
}
}
func TestVaultClassify_DeleteDocLinksByTypes_MultipleTypes(t *testing.T) {
db := testDB(t)
tenantID, agentID := seedTenantAgent(t, db)
ctx := tenantCtx(tenantID)
vs := newVaultStore(db)
tid := tenantID.String()
aid := agentID.String()
// Create source doc and 4 target docs (for 6 different link types + semantic)
docSource := makeVaultDoc(tid, aid, "source.md", "Source")
docs := make([]*store.VaultDocument, 0)
docs = append(docs, docSource)
for i := 0; i < 7; i++ {
doc := makeVaultDoc(tid, aid, "classify/target-"+string(rune('a'+i))+".md", "Target "+string(rune('A'+i)))
if err := vs.UpsertDocument(ctx, doc); err != nil {
t.Fatalf("UpsertDocument target %d: %v", i, err)
}
docs = append(docs, doc)
}
docSource = docs[0]
if err := vs.UpsertDocument(ctx, docSource); err != nil {
t.Fatalf("UpsertDocument source: %v", err)
}
// Create links with 6 classify types + semantic + wikilink (8 total)
linkTypes := []string{
"reference", // classify type 1
"depends_on", // classify type 2
"extends", // classify type 3
"related", // classify type 4
"supersedes", // classify type 5
"contradicts", // classify type 6
"semantic", // legacy auto-classified type
"wikilink", // manual link type
}
for i, lt := range linkTypes {
link := &store.VaultLink{
FromDocID: docSource.ID,
ToDocID: docs[i+1].ID,
LinkType: lt,
Context: "test context for " + lt,
}
if err := vs.CreateLink(ctx, link); err != nil {
t.Fatalf("CreateLink %s: %v", lt, err)
}
}
// Verify: 8 links created
outLinks, err := vs.GetOutLinks(ctx, tid, docSource.ID)
if err != nil {
t.Fatalf("GetOutLinks before delete: %v", err)
}
if len(outLinks) != 8 {
t.Errorf("expected 8 links before delete, got %d", len(outLinks))
}
// Delete all classify types + semantic (7 types), keep wikilink
typesToDelete := []string{
"reference", "depends_on", "extends",
"related", "supersedes", "contradicts",
"semantic",
}
if err := vs.DeleteDocLinksByTypes(ctx, tid, docSource.ID, typesToDelete); err != nil {
t.Fatalf("DeleteDocLinksByTypes: %v", err)
}
// Verify: only wikilink survives
outLinksAfter, err := vs.GetOutLinks(ctx, tid, docSource.ID)
if err != nil {
t.Fatalf("GetOutLinks after delete: %v", err)
}
if len(outLinksAfter) != 1 {
t.Errorf("expected 1 link after deleting 7, got %d", len(outLinksAfter))
}
if outLinksAfter[0].LinkType != "wikilink" {
t.Errorf("expected wikilink, got %s", outLinksAfter[0].LinkType)
}
}
func TestVaultClassify_DeleteDocLinksByTypes_NoMatches(t *testing.T) {
db := testDB(t)
tenantID, agentID := seedTenantAgent(t, db)
ctx := tenantCtx(tenantID)
vs := newVaultStore(db)
tid := tenantID.String()
aid := agentID.String()
// Create docs
docA := makeVaultDoc(tid, aid, "nomatch/a.md", "Doc A")
docB := makeVaultDoc(tid, aid, "nomatch/b.md", "Doc B")
if err := vs.UpsertDocument(ctx, docA); err != nil {
t.Fatalf("UpsertDocument A: %v", err)
}
if err := vs.UpsertDocument(ctx, docB); err != nil {
t.Fatalf("UpsertDocument B: %v", err)
}
// Create link with type "wikilink"
link := &store.VaultLink{
FromDocID: docA.ID,
ToDocID: docB.ID,
LinkType: "wikilink",
Context: "manual link",
}
if err := vs.CreateLink(ctx, link); err != nil {
t.Fatalf("CreateLink: %v", err)
}
// Delete non-existent type — should succeed but delete nothing
if err := vs.DeleteDocLinksByTypes(ctx, tid, docA.ID, []string{"reference"}); err != nil {
t.Fatalf("DeleteDocLinksByTypes non-existent type: %v", err)
}
// Verify wikilink still exists
outLinks, err := vs.GetOutLinks(ctx, tid, docA.ID)
if err != nil {
t.Fatalf("GetOutLinks after no-op delete: %v", err)
}
if len(outLinks) != 1 {
t.Errorf("expected 1 link (unchanged), got %d", len(outLinks))
}
if outLinks[0].LinkType != "wikilink" {
t.Errorf("expected wikilink, got %s", outLinks[0].LinkType)
}
}
func TestVaultClassify_DeleteDocLinksByTypes_EmptyTypeList(t *testing.T) {
db := testDB(t)
tenantID, agentID := seedTenantAgent(t, db)
ctx := tenantCtx(tenantID)
vs := newVaultStore(db)
tid := tenantID.String()
aid := agentID.String()
// Create docs and link
docA := makeVaultDoc(tid, aid, "empty/a.md", "Doc A")
docB := makeVaultDoc(tid, aid, "empty/b.md", "Doc B")
if err := vs.UpsertDocument(ctx, docA); err != nil {
t.Fatalf("UpsertDocument A: %v", err)
}
if err := vs.UpsertDocument(ctx, docB); err != nil {
t.Fatalf("UpsertDocument B: %v", err)
}
link := &store.VaultLink{
FromDocID: docA.ID,
ToDocID: docB.ID,
LinkType: "reference",
Context: "test",
}
if err := vs.CreateLink(ctx, link); err != nil {
t.Fatalf("CreateLink: %v", err)
}
// Delete with empty type list — should succeed but delete nothing
if err := vs.DeleteDocLinksByTypes(ctx, tid, docA.ID, []string{}); err != nil {
t.Fatalf("DeleteDocLinksByTypes empty list: %v", err)
}
// Verify link still exists
outLinks, err := vs.GetOutLinks(ctx, tid, docA.ID)
if err != nil {
t.Fatalf("GetOutLinks after empty delete: %v", err)
}
if len(outLinks) != 1 {
t.Errorf("expected 1 link (unchanged), got %d", len(outLinks))
}
}
func TestVaultClassify_DeleteDocLinksByTypes_MultipleSourceDocs(t *testing.T) {
db := testDB(t)
tenantID, agentID := seedTenantAgent(t, db)
ctx := tenantCtx(tenantID)
vs := newVaultStore(db)
tid := tenantID.String()
aid := agentID.String()
// Create 2 source docs and 1 target doc
docSource1 := makeVaultDoc(tid, aid, "multi/source1.md", "Source 1")
docSource2 := makeVaultDoc(tid, aid, "multi/source2.md", "Source 2")
docTarget := makeVaultDoc(tid, aid, "multi/target.md", "Target")
for _, doc := range []*store.VaultDocument{docSource1, docSource2, docTarget} {
if err := vs.UpsertDocument(ctx, doc); err != nil {
t.Fatalf("UpsertDocument: %v", err)
}
}
// Create links from both sources to target
linkTypes := []string{"reference", "semantic", "wikilink"}
for _, lt := range linkTypes {
link1 := &store.VaultLink{
FromDocID: docSource1.ID,
ToDocID: docTarget.ID,
LinkType: lt,
Context: "from source1",
}
link2 := &store.VaultLink{
FromDocID: docSource2.ID,
ToDocID: docTarget.ID,
LinkType: lt,
Context: "from source2",
}
if err := vs.CreateLink(ctx, link1); err != nil {
t.Fatalf("CreateLink source1: %v", err)
}
if err := vs.CreateLink(ctx, link2); err != nil {
t.Fatalf("CreateLink source2: %v", err)
}
}
// Verify: both sources have 3 links each
outLinks1, err := vs.GetOutLinks(ctx, tid, docSource1.ID)
if err != nil {
t.Fatalf("GetOutLinks source1 before: %v", err)
}
outLinks2, err := vs.GetOutLinks(ctx, tid, docSource2.ID)
if err != nil {
t.Fatalf("GetOutLinks source2 before: %v", err)
}
if len(outLinks1) != 3 || len(outLinks2) != 3 {
t.Errorf("expected 3 links each source before delete, got %d and %d", len(outLinks1), len(outLinks2))
}
// Delete "reference" + "semantic" from source1 only
if err := vs.DeleteDocLinksByTypes(ctx, tid, docSource1.ID, []string{"reference", "semantic"}); err != nil {
t.Fatalf("DeleteDocLinksByTypes source1: %v", err)
}
// Verify: source1 has only "wikilink", source2 unchanged
outLinks1After, err := vs.GetOutLinks(ctx, tid, docSource1.ID)
if err != nil {
t.Fatalf("GetOutLinks source1 after: %v", err)
}
outLinks2After, err := vs.GetOutLinks(ctx, tid, docSource2.ID)
if err != nil {
t.Fatalf("GetOutLinks source2 after: %v", err)
}
if len(outLinks1After) != 1 {
t.Errorf("expected 1 link in source1 after delete, got %d", len(outLinks1After))
}
if outLinks1After[0].LinkType != "wikilink" {
t.Errorf("expected wikilink in source1, got %s", outLinks1After[0].LinkType)
}
if len(outLinks2After) != 3 {
t.Errorf("expected 3 links in source2 (unchanged), got %d", len(outLinks2After))
}
}
func TestVaultClassify_DeleteDocLinksByTypes_TenantIsolation(t *testing.T) {
db := testDB(t)
tenantA, tenantB, agentA, agentB := seedTwoTenants(t, db)
ctxA := tenantCtx(tenantA)
ctxB := tenantCtx(tenantB)
vs := newVaultStore(db)
tidA := tenantA.String()
tidB := tenantB.String()
aidA := agentA.String()
aidB := agentB.String()
// Create docs in tenant A
docA1 := makeVaultDoc(tidA, aidA, "iso/a1.md", "Tenant A Doc 1")
docA2 := makeVaultDoc(tidA, aidA, "iso/a2.md", "Tenant A Doc 2")
if err := vs.UpsertDocument(ctxA, docA1); err != nil {
t.Fatalf("UpsertDocument tenantA doc1: %v", err)
}
if err := vs.UpsertDocument(ctxA, docA2); err != nil {
t.Fatalf("UpsertDocument tenantA doc2: %v", err)
}
// Create docs in tenant B
docB1 := makeVaultDoc(tidB, aidB, "iso/b1.md", "Tenant B Doc 1")
docB2 := makeVaultDoc(tidB, aidB, "iso/b2.md", "Tenant B Doc 2")
if err := vs.UpsertDocument(ctxB, docB1); err != nil {
t.Fatalf("UpsertDocument tenantB doc1: %v", err)
}
if err := vs.UpsertDocument(ctxB, docB2); err != nil {
t.Fatalf("UpsertDocument tenantB doc2: %v", err)
}
// Create links in both tenants
linkA := &store.VaultLink{
FromDocID: docA1.ID,
ToDocID: docA2.ID,
LinkType: "reference",
Context: "tenant A link",
}
linkB := &store.VaultLink{
FromDocID: docB1.ID,
ToDocID: docB2.ID,
LinkType: "reference",
Context: "tenant B link",
}
if err := vs.CreateLink(ctxA, linkA); err != nil {
t.Fatalf("CreateLink tenantA: %v", err)
}
if err := vs.CreateLink(ctxB, linkB); err != nil {
t.Fatalf("CreateLink tenantB: %v", err)
}
// Verify both have links
outA, err := vs.GetOutLinks(ctxA, tidA, docA1.ID)
if err != nil {
t.Fatalf("GetOutLinks tenantA before: %v", err)
}
outB, err := vs.GetOutLinks(ctxB, tidB, docB1.ID)
if err != nil {
t.Fatalf("GetOutLinks tenantB before: %v", err)
}
if len(outA) != 1 || len(outB) != 1 {
t.Errorf("expected 1 link each before delete, got %d and %d", len(outA), len(outB))
}
// Delete from tenant A only — should NOT affect tenant B
if err := vs.DeleteDocLinksByTypes(ctxA, tidA, docA1.ID, []string{"reference"}); err != nil {
t.Fatalf("DeleteDocLinksByTypes tenantA: %v", err)
}
// Verify: tenantA link deleted, tenantB link unchanged
outAAfter, err := vs.GetOutLinks(ctxA, tidA, docA1.ID)
if err != nil {
t.Fatalf("GetOutLinks tenantA after: %v", err)
}
outBAfter, err := vs.GetOutLinks(ctxB, tidB, docB1.ID)
if err != nil {
t.Fatalf("GetOutLinks tenantB after: %v", err)
}
if len(outAAfter) != 0 {
t.Errorf("expected 0 links in tenantA after delete, got %d", len(outAAfter))
}
if len(outBAfter) != 1 {
t.Errorf("expected 1 link in tenantB (unchanged), got %d", len(outBAfter))
}
}
func TestVaultClassify_DeleteDocLinksByTypes_NonExistentDoc(t *testing.T) {
db := testDB(t)
tenantID, _ := seedTenantAgent(t, db)
ctx := tenantCtx(tenantID)
vs := newVaultStore(db)
tid := tenantID.String()
// Try to delete links from non-existent doc — should succeed (no-op)
fakeDocID := "fake-doc-id-12345"
if err := vs.DeleteDocLinksByTypes(ctx, tid, fakeDocID, []string{"reference"}); err != nil {
t.Fatalf("DeleteDocLinksByTypes on non-existent doc: %v", err)
}
// GetOutLinks for non-existent doc should return empty or error gracefully
outLinks, err := vs.GetOutLinks(ctx, tid, fakeDocID)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
// Some implementations may error on non-existent doc, others return empty list
// Both are acceptable for a non-existent doc
}
if outLinks != nil && len(outLinks) != 0 {
t.Errorf("expected empty or error for non-existent doc, got %d links", len(outLinks))
}
}
func TestVaultClassify_DeleteAllClassifyTypesAndSemantic(t *testing.T) {
db := testDB(t)
tenantID, agentID := seedTenantAgent(t, db)
ctx := tenantCtx(tenantID)
vs := newVaultStore(db)
tid := tenantID.String()
aid := agentID.String()
// Create documents
docSource := makeVaultDoc(tid, aid, "classify/source.md", "Source Doc")
docTarget1 := makeVaultDoc(tid, aid, "classify/target1.md", "Target 1")
docTarget2 := makeVaultDoc(tid, aid, "classify/target2.md", "Target 2")
for _, doc := range []*store.VaultDocument{docSource, docTarget1, docTarget2} {
if err := vs.UpsertDocument(ctx, doc); err != nil {
t.Fatalf("UpsertDocument: %v", err)
}
}
// Create links with mix of classify types, semantic, and manual types
linkConfigs := []struct {
toDocID string
linkType string
}{
{docTarget1.ID, "reference"},
{docTarget1.ID, "semantic"},
{docTarget2.ID, "depends_on"},
{docTarget2.ID, "wikilink"},
}
for _, cfg := range linkConfigs {
link := &store.VaultLink{
FromDocID: docSource.ID,
ToDocID: cfg.toDocID,
LinkType: cfg.linkType,
Context: "test context",
}
if err := vs.CreateLink(ctx, link); err != nil {
t.Fatalf("CreateLink %s: %v", cfg.linkType, err)
}
}
// Verify: 4 links created
outLinks, err := vs.GetOutLinks(ctx, tid, docSource.ID)
if err != nil {
t.Fatalf("GetOutLinks before delete: %v", err)
}
if len(outLinks) != 4 {
t.Errorf("expected 4 links before delete, got %d", len(outLinks))
}
// Delete all classify types (reference, depends_on) + semantic, keep wikilink
if err := vs.DeleteDocLinksByTypes(ctx, tid, docSource.ID, []string{
"reference", "depends_on", "extends", "related", "supersedes", "contradicts", "semantic",
}); err != nil {
t.Fatalf("DeleteDocLinksByTypes: %v", err)
}
// Verify: only wikilink remains
outLinksAfter, err := vs.GetOutLinks(ctx, tid, docSource.ID)
if err != nil {
t.Fatalf("GetOutLinks after delete: %v", err)
}
if len(outLinksAfter) != 1 {
t.Errorf("expected 1 link after cleanup, got %d", len(outLinksAfter))
}
if outLinksAfter[0].LinkType != "wikilink" {
t.Errorf("expected wikilink survivor, got %s", outLinksAfter[0].LinkType)
}
}
func TestVaultClassify_DeleteDocLinksByTypes_CasePreservation(t *testing.T) {
db := testDB(t)
tenantID, agentID := seedTenantAgent(t, db)
ctx := tenantCtx(tenantID)
vs := newVaultStore(db)
tid := tenantID.String()
aid := agentID.String()
// Create documents
docA := makeVaultDoc(tid, aid, "case/a.md", "Doc A")
docB := makeVaultDoc(tid, aid, "case/b.md", "Doc B")
if err := vs.UpsertDocument(ctx, docA); err != nil {
t.Fatalf("UpsertDocument A: %v", err)
}
if err := vs.UpsertDocument(ctx, docB); err != nil {
t.Fatalf("UpsertDocument B: %v", err)
}
// Create links with different case sensitivity
link := &store.VaultLink{
FromDocID: docA.ID,
ToDocID: docB.ID,
LinkType: "depends_on", // lowercase with underscore
Context: "test",
}
if err := vs.CreateLink(ctx, link); err != nil {
t.Fatalf("CreateLink: %v", err)
}
// Try deleting with exact case match
if err := vs.DeleteDocLinksByTypes(ctx, tid, docA.ID, []string{"depends_on"}); err != nil {
t.Fatalf("DeleteDocLinksByTypes: %v", err)
}
// Verify link is deleted
outLinks, err := vs.GetOutLinks(ctx, tid, docA.ID)
if err != nil {
t.Fatalf("GetOutLinks: %v", err)
}
if len(outLinks) != 0 {
t.Errorf("expected 0 links after exact case match delete, got %d", len(outLinks))
}
}