diff --git a/internal/ingestion/knowledge_compile/consumer.go b/internal/ingestion/knowledge_compile/consumer.go index 9e6e3d4a98..e0462188e4 100644 --- a/internal/ingestion/knowledge_compile/consumer.go +++ b/internal/ingestion/knowledge_compile/consumer.go @@ -1144,7 +1144,20 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, e Score: g.score, }) } - batched, err = deduper.DecideBatch(ctx, batched) + llmCtx, finishLLMProgress := withKnowledgeCompileLLMProgress(ctx, func(progressCtx context.Context, component, message string) { + c.reportProgress( + progressCtx, + tenant, + kb, + token, + taskTypes, + 0.82, + "llm_merge", + fmt.Sprintf("%s: %s", component, message), + ) + }) + batched, err = deduper.DecideBatch(llmCtx, batched) + finishLLMProgress() if err != nil { return err } diff --git a/internal/ingestion/knowledge_compile/stub.go b/internal/ingestion/knowledge_compile/stub.go new file mode 100644 index 0000000000..f58d15bcaa --- /dev/null +++ b/internal/ingestion/knowledge_compile/stub.go @@ -0,0 +1,12 @@ +//go:build !ee + +package knowledge_compile + +import "context" + +func withKnowledgeCompileLLMProgress( + ctx context.Context, + _ func(context.Context, string, string), +) (context.Context, func()) { + return ctx, func() {} +} diff --git a/internal/ingestion/task/knowledge_compiler_wiring.go b/internal/ingestion/task/knowledge_compiler_wiring.go index af98517bd5..365c86fe6b 100644 --- a/internal/ingestion/task/knowledge_compiler_wiring.go +++ b/internal/ingestion/task/knowledge_compiler_wiring.go @@ -602,7 +602,7 @@ func newKnowledgeCompilerDepsResolver() kc.DepsResolver { }) return kc.Deps{ - Chat: &kcChatInvoker{factory: modelFactory, tenantID: tenantID, llmID: llmID}, + Chat: &kcChatInvoker{factory: modelFactory, tenantID: tenantID, llmID: llmID, llmProgress: newKCLLMProgressReporter()}, Embed: embedder, WikiPages: &kcWikiPageStore{docEngine: kcWiringDeps.Engine}, WikiMapVersions: mapStore, @@ -617,9 +617,10 @@ func newKnowledgeCompilerDepsResolver() kc.DepsResolver { // kcChatInvoker adapts a factory-created ChatModel to the knowledge compiler. type kcChatInvoker struct { - factory *service.ModelFactory - tenantID string - llmID string + factory *service.ModelFactory + tenantID string + llmID string + llmProgress kcLLMProgressReporter } func (c *kcChatInvoker) Chat(ctx context.Context, req kc.ChatRequest) (*kc.ChatResponse, error) { @@ -629,6 +630,9 @@ func (c *kcChatInvoker) Chat(ctx context.Context, req kc.ChatRequest) (*kc.ChatR } chatModel, err := c.factory.NewChatModel(permission.WithTenantID(ctx, c.tenantID), llmID) if err != nil { + if !req.DisableRetry && c.llmProgress != nil { + c.llmProgress.report(ctx, "[ERROR] LLM call failed: "+kc.CompactError(err)) + } return nil, err } msgs := []models.Message{ @@ -685,7 +689,13 @@ func (c *kcChatInvoker) Chat(ctx context.Context, req kc.ChatRequest) (*kc.ChatR // providers legitimately need several minutes to return a response. attemptCtx, cancel := context.WithTimeout(ctx, kcChatAttemptTimeout) defer cancel() - r, err := chatModel.ChatWithMessages(attemptCtx, msgs, config, nil) + r, err := func() (*models.ChatResponse, error) { + if c.llmProgress == nil { + return chatModel.ChatWithMessages(attemptCtx, msgs, config, nil) + } + defer c.llmProgress.begin(ctx)() + return chatModel.ChatWithMessages(attemptCtx, msgs, config, nil) + }() if err != nil { if !req.DisableRetry { final := !appcommon.IsTransientError(err) || attempt > kcChatRetryMax diff --git a/internal/ingestion/task/stub.go b/internal/ingestion/task/stub.go new file mode 100644 index 0000000000..be4b4d5f31 --- /dev/null +++ b/internal/ingestion/task/stub.go @@ -0,0 +1,22 @@ +//go:build !ee + +package task + +import "context" + +type kcLLMProgressReporter interface { + begin(ctx context.Context) func() + report(ctx context.Context, message string) +} + +type noopKCLLMProgressReporter struct{} + +func newKCLLMProgressReporter() kcLLMProgressReporter { + return noopKCLLMProgressReporter{} +} + +func (noopKCLLMProgressReporter) begin(context.Context) func() { + return func() {} +} + +func (noopKCLLMProgressReporter) report(context.Context, string) {}