fix(vault): increase max_tokens + hot-reload provider for enrichment

- Increase classify max_tokens 1024→2048, summarize 1536→4096 to
  prevent truncated JSON from models like gemini-2.5-pro
- Add debug logging: raw LLM output on parse failures, finish_reason
  truncation warnings with model name
- Hot-swap vault enrichment provider/model on config change without
  restart (wired into TopicSystemConfigChanged handler)
- Use RWMutex-guarded llm() accessor for thread-safe provider reads
This commit is contained in:
viettranx committed 2026-04-12 15:08:59 +07:00
1 parent fbfae1e618
commit ba66a7b67d
5 files changed
+56 -14

No files matched your search

+5 -2
View File
@@ -212,8 +212,9 @@ func runGateway() {
// V3: Wire vault enrichment worker (async summary + embedding + auto-linking).
var enrichProgress *vault.EnrichProgress
var updateVaultProvider vault.ProviderUpdater
if pgStores.Vault != nil && bgProvider != nil {
cleanupVaultEnrich, ep := vault.RegisterEnrichWorker(vault.EnrichWorkerDeps{
cleanupVaultEnrich, ep, updater := vault.RegisterEnrichWorker(vault.EnrichWorkerDeps{
VaultStore: pgStores.Vault,
Provider: bgProvider,
Model: bgModel,
@@ -222,6 +223,7 @@ func runGateway() {
TeamStore: pgStores.Teams, // Phase 04 task-based auto-linking
})
enrichProgress = ep
updateVaultProvider = updater
defer cleanupVaultEnrich()
slog.Info("vault enrichment worker registered", "provider", bgProvider.Name(), "model", bgModel)
}
@@ -296,7 +298,8 @@ func runGateway() {
agentRouter: agentRouter,
toolsReg: toolsReg,
skillsLoader: skillsLoader,
enrichProgress: enrichProgress,
enrichProgress: enrichProgress,
updateVaultProvider: updateVaultProvider,
workspace: workspace,
dataDir: dataDir,
domainBus: domainBus,
+2 -1
View File
@@ -28,7 +28,8 @@ type gatewayDeps struct {
toolsReg *tools.Registry
skillsLoader *skills.Loader // optional: enables skill creation in evolution approval
permCache *cache.PermissionCache // nil if no tenant store; closed on shutdown to stop sweep goroutines
enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered
enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered
updateVaultProvider vault.ProviderUpdater // nil if enrichment worker not registered; hot-swaps LLM
workspace string
dataDir string
domainBus eventbus.DomainEventBus
+6
View File
@@ -118,6 +118,12 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer(
pgMem.UpdateChunkConfig(mem.MaxChunkLen, mem.ChunkOverlap)
}
}
// Hot-swap vault enrichment provider/model if config changed
if d.updateVaultProvider != nil {
if p, m := resolveBackgroundProvider(d.cfg, d.providerRegistry); p != nil {
d.updateVaultProvider(p, m)
}
}
slog.Debug("system_configs refreshed to in-memory config", "keys", len(sysConfigs))
}
})
+7 -4
View File
@@ -12,7 +12,7 @@ import (
)
const (
classifyMaxTokens = 1024
classifyMaxTokens = 2048
classifyTemperature = 0.1
classifyCtxMaxLen = 256 // max context string length stored in DB
classifySummaryMaxChars = 300 // max summary chars in prompt (validated: 300 for accuracy)
@@ -37,7 +37,8 @@ type candidatePair struct {
// 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 {
provider, _ := w.llm()
if provider == nil {
return
}
@@ -77,6 +78,7 @@ func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID stri
parsed, err := parseClassifyResponse(raw, len(chunk))
if err != nil {
slog.Warn("vault.classify: parse_failed_first", "doc", sourceDocID, "err", err, "raw_len", len(raw), "raw", raw)
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 {
@@ -85,7 +87,7 @@ func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID stri
}
parsed, err = parseClassifyResponse(raw2, len(chunk))
if err != nil {
slog.Warn("vault.classify: parse_still_failed", "doc", sourceDocID, "err", err)
slog.Warn("vault.classify: parse_still_failed", "doc", sourceDocID, "err", err, "raw_len", len(raw2), "raw", raw2)
continue
}
}
@@ -176,12 +178,13 @@ func (w *enrichWorker) gatherCandidates(ctx context.Context, tenantID, _ string,
// callClassifyWithRetry calls the LLM with shared retry logic.
func (w *enrichWorker) callClassifyWithRetry(ctx context.Context, system, user string) (string, error) {
_, model := w.llm()
return w.chatWithRetry(ctx, "vault.classify", providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: system},
{Role: "user", Content: user},
},
Model: w.model,
Model: model,
Options: map[string]any{"max_tokens": classifyMaxTokens, "temperature": classifyTemperature},
})
}
+36 -7
View File
@@ -45,9 +45,12 @@ type EnrichWorkerDeps struct {
TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe)
}
// ProviderUpdater allows hot-swapping the LLM provider/model at runtime.
type ProviderUpdater func(p providers.Provider, model string)
// RegisterEnrichWorker subscribes the enrichment worker to vault doc events.
// Returns (unsubscribe func, progress tracker for WS broadcast).
func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress) {
// Returns (unsubscribe func, progress tracker, provider updater for hot-reload).
func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, ProviderUpdater) {
progress := NewEnrichProgress(deps.MsgBus)
w := &enrichWorker{
vault: deps.VaultStore,
@@ -59,7 +62,7 @@ func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress) {
progress: progress,
}
unsub := deps.EventBus.Subscribe(eventbus.EventVaultDocUpserted, w.Handle)
return unsub, progress
return unsub, progress, w.UpdateProvider
}
// enrichWorker processes vault document upsert events to generate summaries,
@@ -69,6 +72,7 @@ type enrichWorker struct {
teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil
provider providers.Provider
model string
llmMu sync.RWMutex // guards provider + model hot-swap
queue enrichBatchQueue
progress *EnrichProgress
@@ -78,6 +82,26 @@ type enrichWorker struct {
sem *semaphore.Weighted // limits concurrent LLM summarize calls
}
// UpdateProvider hot-swaps the LLM provider and model used by the enrichment worker.
// Called when background worker config changes at runtime.
func (w *enrichWorker) UpdateProvider(p providers.Provider, model string) {
w.llmMu.Lock()
defer w.llmMu.Unlock()
if w.provider != nil && w.provider.Name() == p.Name() && w.model == model {
return // no change
}
w.provider = p
w.model = model
slog.Info("vault.enrich: provider updated", "provider", p.Name(), "model", model)
}
// llm returns the current provider and model, safe for concurrent reads.
func (w *enrichWorker) llm() (providers.Provider, string) {
w.llmMu.RLock()
defer w.llmMu.RUnlock()
return w.provider, w.model
}
// enrichTaskSiblingCap bounds the number of auto-linked siblings per
// (source_doc × task) pair. Tunable via VAULT_TASK_SIBLING_CAP env var so
// operators can raise/lower without a rebuild.
@@ -337,13 +361,14 @@ func (w *enrichWorker) batchSummarize(ctx context.Context, paths, contents []str
fmt.Fprintf(&b, "[%d] File: %s\n%s\n\n", i+1, paths[i], contents[i])
}
_, model := w.llm()
raw, err := w.chatWithRetry(ctx, "vault.batch_summarize", providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: batchSummarizePrompt},
{Role: "user", Content: b.String()},
},
Model: w.model,
Options: map[string]any{"max_tokens": 1536, "temperature": 0.2},
Model: model,
Options: map[string]any{"max_tokens": 4096, "temperature": 0.2},
})
if err != nil {
slog.Warn("vault.enrich: batch_summarize", "count", len(paths), "err", err)
@@ -371,7 +396,7 @@ func parseBatchSummaries(raw string, expected int) []string {
Summary string `json:"summary"`
}
if err := json.Unmarshal([]byte(raw), &results); err != nil {
slog.Warn("vault.enrich: parse_batch_summaries", "err", err)
slog.Warn("vault.enrich: parse_batch_summaries", "err", err, "raw_len", len(raw), "raw", raw)
return nil
}
@@ -397,14 +422,18 @@ func (w *enrichWorker) chatWithRetry(ctx context.Context, logPrefix string, req
case <-time.After(enrichRetryBackoffs[attempt]):
}
}
provider, _ := w.llm()
cctx, cancel := context.WithTimeout(ctx, enrichRetryTimeouts[attempt])
resp, err := w.provider.Chat(cctx, req)
resp, err := provider.Chat(cctx, req)
cancel()
if err != nil {
lastErr = err
slog.Warn(logPrefix+": retry", "attempt", attempt+1, "err", err)
continue
}
if resp.FinishReason == "length" {
slog.Warn(logPrefix+": truncated", "finish_reason", "length", "model", req.Model, "content_len", len(resp.Content), "max_tokens", req.Options["max_tokens"])
}
return strings.TrimSpace(resp.Content), nil
}
return "", fmt.Errorf("%s exhausted %d retries: %w", logPrefix, enrichMaxRetries, lastErr)