mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-07-13 05:05:52 +00:00
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
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -145,6 +146,30 @@ func (s *PGVaultStore) DeleteDocLinksByType(ctx context.Context, tenantID, docID
|
||||
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) {
|
||||
var links []store.VaultLink
|
||||
for rows.Next() {
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -122,6 +123,29 @@ func (s *SQLiteVaultStore) DeleteDocLinksByType(ctx context.Context, tenantID, d
|
||||
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) {
|
||||
var links []store.VaultLink
|
||||
for rows.Next() {
|
||||
|
||||
@@ -93,6 +93,7 @@ type VaultStore interface {
|
||||
GetBacklinks(ctx context.Context, tenantID, docID string) ([]VaultBacklink, error)
|
||||
DeleteDocLinks(ctx context.Context, tenantID, docID string) error
|
||||
DeleteDocLinksByType(ctx context.Context, tenantID, docID, linkType string) error
|
||||
DeleteDocLinksByTypes(ctx context.Context, tenantID, docID string, types []string) error
|
||||
|
||||
// Enrichment
|
||||
// UpdateSummaryAndReembed updates summary text and re-generates embedding from title+path+summary.
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,6 @@ package vault
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -80,6 +79,12 @@ func (w *enrichWorker) Handle(ctx context.Context, event eventbus.DomainEvent) e
|
||||
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.
|
||||
func (w *enrichWorker) processBatch(ctx context.Context, key string) {
|
||||
for {
|
||||
@@ -91,10 +96,6 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) {
|
||||
continue
|
||||
}
|
||||
|
||||
type enriched struct {
|
||||
payload eventbus.VaultDocUpsertedPayload
|
||||
summary string
|
||||
}
|
||||
var results []enriched
|
||||
|
||||
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})
|
||||
}
|
||||
|
||||
// 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 {
|
||||
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)
|
||||
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)
|
||||
|
||||
// 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)
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
func (w *enrichWorker) recordDedup(docID, hash string) {
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user