mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-11 12:18:59 +00:00
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:
1 parent
fbfae1e618
commit
ba66a7b67d
5 files changed
+56
-14
No files matched your search
+5
-2
@@ -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
@@ -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
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
})
|
||||
|
||||
@@ -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},
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user