From 86b6f64abb8e241c1f8bc1a450d318b984c06a01 Mon Sep 17 00:00:00 2001 From: jay77721 <164177721+jay77721@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:45:01 +0800 Subject: [PATCH] fix(ingestion): fit extractor and tagger prompts to model context (#18095) Trim Extractor call prompts and the automatic tagger prompt to the chat model's context window (`content_length`) before sending, so oversized chunks or tag files are trimmed instead of rejected by the provider with a context-length error. --- internal/ingestion/component/extractor.go | 159 +++++++++++ internal/ingestion/component/extractor_tag.go | 14 +- .../ingestion/component/extractor_tag_test.go | 81 ++++++ .../ingestion/component/extractor_test.go | 267 +++++++++++++++++- 4 files changed, 519 insertions(+), 2 deletions(-) diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 177b31f8f1..f9f1295644 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -80,6 +80,8 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/common" + "ragflow/internal/component/messagefit" + "ragflow/internal/dao" "ragflow/internal/engine/redis" "ragflow/internal/entity" "ragflow/internal/entity/models" @@ -1134,6 +1136,11 @@ func (c *ExtractorComponent) callRaw(ctx context.Context, db *gorm.DB, in extrac return nil, err } msgs := buildExtractorMessages(in.systemPrompt, in.prompt, chunkText, in.chunks) + fitted, fitErr := fitExtractorMessages(ctx, db, in.llmID, msgs) + if fitErr != nil { + return nil, fitErr + } + msgs = fitted inv := getExtractorChatInvoker() req := extractorChatRequest{ Driver: driver, @@ -1354,6 +1361,158 @@ func isBareTenantModelID(s string) bool { return true } +// extractorContextLengthOverride is a narrow test seam mirroring +// extractorChatTargetResolverOverride: it lets unit tests supply a context +// length without a real tenant model row, so the message-fitting wiring in +// callRaw/llmTagChunk can be exercised without a DB. When set, +// extractorContextLength consults it first. +var ( + extractorContextLengthOverrideMu sync.RWMutex + extractorContextLengthOverride func(ctx context.Context, llmID string) int +) + +// SetExtractorContextLengthOverride swaps the package-level context-length +// resolver for tests. Pass nil to restore the default. Concurrent-safe. +func SetExtractorContextLengthOverride(fn func(ctx context.Context, llmID string) int) { + extractorContextLengthOverrideMu.Lock() + defer extractorContextLengthOverrideMu.Unlock() + extractorContextLengthOverride = fn +} + +func getExtractorContextLengthOverride() func(ctx context.Context, llmID string) int { + extractorContextLengthOverrideMu.RLock() + defer extractorContextLengthOverrideMu.RUnlock() + return extractorContextLengthOverride +} + +// extractorContextLength returns the chat model's context window +// (content_length) for the effective chat model used by a call, or 0 when +// unavailable. Mirrors Python's chat_mdl.max_length. Used as the token +// budget for message fitting so oversized prompts are trimmed instead of +// rejected by the provider. When llm_id is empty the call falls back to the +// tenant default chat model (see resolveExtractorChatTarget), so the same +// model is resolved here; otherwise the default-model path would never get +// message fitting. Returns 0 (skip fitting) when the model is unknown (e.g. +// unit tests with synthetic llm_id or no canvas state). +func extractorContextLength(ctx context.Context, db *gorm.DB, llmID string) int { + if fn := getExtractorContextLengthOverride(); fn != nil { + return fn(ctx, llmID) + } + if db == nil { + db = dao.DB + } + state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + if err != nil || state == nil { + return 0 + } + tidVal, _ := state.GetGlobal("tenant_id") + tid, _ := tidVal.(string) + if tid == "" { + return 0 + } + if llmID == "" { + llmID = defaultChatModelRef(ctx, db, tid) + } + if llmID == "" { + return 0 + } + return dao.ResolveModelContentLength(ctx, db, llmID, "", "") +} + +// defaultChatModelRef returns the tenant's default chat model reference — +// the tenant_model UUID when one is pinned, otherwise the composite +// "model@provider" id — or "" when the tenant has no default chat model. +func defaultChatModelRef(ctx context.Context, db *gorm.DB, tenantID string) string { + if db == nil { + // No database to read the tenant's default chat model from. + return "" + } + tenant, err := dao.NewTenantDAO().GetByID(ctx, db, tenantID) + if err != nil || tenant == nil { + return "" + } + if tenant.TenantLLMID != nil && *tenant.TenantLLMID != "" { + return *tenant.TenantLLMID + } + return tenant.LLMID +} + +// extractorContextFitBudget returns 97% of the model's context window as the +// fitting budget, mirroring the agent component's contextFitBudget. The +// margin leaves headroom for the difference between the cl100k tokenizer used +// for counting and the model's own tokenizer, plus per-message formatting +// overhead, so a fitted prompt stays inside the provider's real context limit +// instead of landing exactly on it. +func extractorContextFitBudget(ctxLen int) int { + budget := int(float64(ctxLen) * 0.97) + if budget < 1 { + // Never hand messagefit a <=0 budget: Fit treats <=0 as the 8192 + // default, which would stop trimming entirely for a tiny context. + return 1 + } + return budget +} + +// fitExtractorMessages trims msgs to the chat model's context window using +// the shared messagefit fitter (mirrors Python's message_fit_in), dropping +// entries the fitter removed. It returns a clear error instead of letting a +// conversation whose final user turn was trimmed to empty reach the provider: +// the proportional trim can do that when the system prompt alone exceeds the +// context budget, and providers reject empty user turns with an obscure error +// after retries. +func fitExtractorMessages(ctx context.Context, db *gorm.DB, llmID string, msgs []eschema.Message) ([]eschema.Message, error) { + ctxLen := extractorContextLength(ctx, db, llmID) + if ctxLen <= 0 { + return msgs, nil + } + fitMsgs := make([]messagefit.Message, len(msgs)) + for i := range msgs { + fitMsgs[i] = messagefit.Message{Role: string(msgs[i].Role), Content: msgs[i].Content} + } + kept, keptIdx, _ := messagefit.Fit(fitMsgs, extractorContextFitBudget(ctxLen)) + + fitted := make([]eschema.Message, 0, len(kept)) + for j, i := range keptIdx { + msgs[i].Content = kept[j].Content + fitted = append(fitted, msgs[i]) + } + if len(fitted) == 0 { + return nil, errors.New("extractor: message fitting dropped every message; check the chat model context length setting") + } + // The system prompt carries the extraction contract (output format, + // field definitions); sending without it would silently produce + // garbage. The proportional trim can empty every system message when + // the final user turn alone fills the budget, so require at least one + // retained system message with non-empty content — a system message kept + // but trimmed to empty is just as useless as a dropped one. The guard only + // applies when the input actually had a system message: systemPrompt is + // optional and a user-only prompt is a valid request. + hadSystem := false + for _, m := range msgs { + if m.Role == eschema.System { + hadSystem = true + break + } + } + if hadSystem { + keptSystem := false + for _, m := range fitted { + if m.Role == eschema.System && strings.TrimSpace(m.Content) != "" { + keptSystem = true + break + } + } + if !keptSystem { + return nil, errors.New("extractor: message fitting emptied the system prompt; check the chat model context length setting or reduce the prompt size") + } + } + last := fitted[len(fitted)-1] + if last.Role != eschema.User || strings.TrimSpace(last.Content) == "" { + return nil, errors.New("extractor: message fitting emptied the final user turn; check the chat model context length setting or reduce the prompt size") + } + return fitted, nil +} + // buildExtractorMessages assembles system + user messages for // one extraction call. The user prompt is rendered as // "\n\n" so the python behavior of diff --git a/internal/ingestion/component/extractor_tag.go b/internal/ingestion/component/extractor_tag.go index 9a2b0590a7..6d4dad1235 100644 --- a/internal/ingestion/component/extractor_tag.go +++ b/internal/ingestion/component/extractor_tag.go @@ -191,7 +191,7 @@ func (c *ExtractorComponent) runAutoTags(ctx context.Context, db *gorm.DB, in ex case <-ctx.Done(): return } - llmTagChunk(ctx, inv, docsToTag[idx], indexed.allTags, examples, in.llmID, driver, model, apiKey, baseURL, topN) + llmTagChunk(ctx, db, inv, docsToTag[idx], indexed.allTags, examples, in.llmID, driver, model, apiKey, baseURL, topN) }(i) } wg.Wait() @@ -652,6 +652,7 @@ func roundInt(f float64) int { func llmTagChunk( ctx context.Context, + db *gorm.DB, inv extractorChatInvoker, chunk map[string]any, allTags map[string]float64, @@ -687,6 +688,17 @@ func llmTagChunk( {Role: eschema.System, Content: prompt}, {Role: eschema.User, Content: "Output:"}, } + // Trim the prompt to the model's context window before sending. The + // system prompt embeds the full chunk text, the entire tag set and up + // to two full examples, so oversized chunks or tag files would + // otherwise be rejected by the provider (context length exceeded). + // Mirrors Python's message_fit_in in content_tagging (generator.py:331). + fitted, fitErr := fitExtractorMessages(ctx, db, llmID, msgs) + if fitErr != nil { + common.Warn("extractor tags: skipping LLM tagging, message fitting failed", zap.Error(fitErr)) + return + } + msgs = fitted temperature := 0.5 var result map[string]int diff --git a/internal/ingestion/component/extractor_tag_test.go b/internal/ingestion/component/extractor_tag_test.go index 6161281e1a..9e7c8df374 100644 --- a/internal/ingestion/component/extractor_tag_test.go +++ b/internal/ingestion/component/extractor_tag_test.go @@ -12,6 +12,7 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/common" "ragflow/internal/ingestion/component/schema" + "ragflow/internal/tokenizer" ) type stubExtractorTagChat struct { @@ -236,6 +237,86 @@ func stubXLSXBytes(t *testing.T) []byte { return buf.Bytes() } +// capturingExtractorTagChat records the request it was given so tests can +// assert on the exact messages sent to the LLM. +type capturingExtractorTagChat struct { + req extractorChatRequest +} + +func (c *capturingExtractorTagChat) Chat(_ context.Context, req extractorChatRequest) (*extractorChatResponse, error) { + c.req = req + return &extractorChatResponse{Content: `{"RAG": 8, "vector database": 6}`}, nil +} + +func pushCapturingTagChat(t *testing.T) *capturingExtractorTagChat { + t.Helper() + capt := &capturingExtractorTagChat{} + SetExtractorChatInvoker(capt) + t.Cleanup(func() { SetExtractorChatInvoker(nil) }) + return capt +} + +func longChunkText() string { + return strings.Repeat("RAGFlow is an open source retrieval augmented generation engine. ", 10) +} + +// TestLlmtagChunk_MessageFit verifies that llmTagChunk trims the prompt to the +// model's context window before sending, mirroring Python's message_fit_in in +// content_tagging (generator.py:331). The tagger system prompt embeds the full +// chunk text plus the whole tag set, so an oversized chunk must be trimmed +// rather than rejected by the provider with "context length exceeded". +func TestLlmtagChunk_MessageFit(t *testing.T) { + const budget = 20 + SetExtractorContextLengthOverride(func(_ context.Context, _ string) int { return budget }) + t.Cleanup(func() { SetExtractorContextLengthOverride(nil) }) + + capt := pushCapturingTagChat(t) + + longText := longChunkText() + chunk := map[string]any{"content_with_weight": longText} + allTags := map[string]float64{"RAG": 1, "database": 1, "AI": 1} + examples := []schema.TaggedChunk{{Content: "example one", TagWeights: map[string]int{"AI": 5}}} + + llmTagChunk(t.Context(), nil, capt, chunk, allTags, examples, "test@test", "test_driver", "test_model", "test_key", "", 3) + + if len(capt.req.Messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(capt.req.Messages)) + } + // The full chunk text must have been trimmed away from the system prompt. + if strings.Contains(capt.req.Messages[0].Content, longText) { + t.Fatal("system prompt was not trimmed to the context budget") + } + total := tokenizer.NumTokensFromString(capt.req.Messages[0].Content) + + tokenizer.NumTokensFromString(capt.req.Messages[1].Content) + // The budget is 97% of the resolved content_length (the override value), + // mirroring the agent's contextFitBudget; the margin keeps a fitted + // prompt inside the provider's real context limit. + if total > extractorContextFitBudget(budget) { + t.Fatalf("fitted messages total %d tokens exceeds the margin-adjusted budget %d", total, extractorContextFitBudget(budget)) + } +} + +// TestLlmtagChunk_NoContextLength_SkipsFit verifies the guard: when the model's +// context length is unknown (extractorContextLength returns 0), the tagger +// passes the prompt through untrimmed instead of erroring. +func TestLlmtagChunk_NoContextLength_SkipsFit(t *testing.T) { + capt := pushCapturingTagChat(t) + + longText := longChunkText() + chunk := map[string]any{"content_with_weight": longText} + allTags := map[string]float64{"RAG": 1} + examples := []schema.TaggedChunk{{Content: "example", TagWeights: map[string]int{"AI": 5}}} + + llmTagChunk(t.Context(), nil, capt, chunk, allTags, examples, "test@test", "test_driver", "test_model", "test_key", "", 3) + + if len(capt.req.Messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(capt.req.Messages)) + } + if !strings.Contains(capt.req.Messages[0].Content, longText) { + t.Fatal("system prompt should pass through untrimmed when context length is unknown") + } +} + func TestParseCSVQuoteAwareReader(t *testing.T) { // A quoted content field containing a comma must not be split into extra // columns: "RAGFlow, the guide",RAG is two fields, not three. diff --git a/internal/ingestion/component/extractor_test.go b/internal/ingestion/component/extractor_test.go index d10236d2b2..0e882b23df 100644 --- a/internal/ingestion/component/extractor_test.go +++ b/internal/ingestion/component/extractor_test.go @@ -19,7 +19,6 @@ package component import ( "context" "errors" - "ragflow/internal/dao" "strings" "sync" "sync/atomic" @@ -27,10 +26,15 @@ import ( "time" eschema "github.com/cloudwego/eino/schema" + "github.com/glebarez/sqlite" + "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" "ragflow/internal/ingestion/component/schema" + "ragflow/internal/tokenizer" "ragflow/internal/utility" ) @@ -1792,3 +1796,264 @@ func TestExtractorComponent_Invoke_FieldValueContainsPlaceholderSubstring(t *tes t.Errorf("expected title substitution to produce literal '{text}' label: %q", userContent) } } + +// TestFitExtractorMessages_RejectsEmptyUserTurn verifies that when +// messagefit's proportional trim would empty the final user turn (the system +// prompt alone exceeds the context budget), the extractor surfaces a clear +// error instead of sending [system, user:""] to the provider. +func TestFitExtractorMessages_RejectsEmptyUserTurn(t *testing.T) { + SetExtractorContextLengthOverride(func(_ context.Context, _ string) int { return 500 }) + t.Cleanup(func() { SetExtractorContextLengthOverride(nil) }) + + msgs := []eschema.Message{ + {Role: eschema.System, Content: strings.Repeat("s ", 1000)}, + {Role: eschema.User, Content: strings.Repeat("u ", 400)}, + } + if _, err := fitExtractorMessages(t.Context(), nil, "test@test", msgs); err == nil { + t.Fatal("expected an error when fitting empties the user turn") + } +} + +// TestFitExtractorMessages_KeepsUserTurn verifies the happy path: with a +// normal budget the fitter trims oversized prompts and the final user turn +// survives, so no error is returned. +func TestFitExtractorMessages_KeepsUserTurn(t *testing.T) { + SetExtractorContextLengthOverride(func(_ context.Context, _ string) int { return 2000 }) + t.Cleanup(func() { SetExtractorContextLengthOverride(nil) }) + + msgs := []eschema.Message{ + {Role: eschema.System, Content: "you are a helpful assistant"}, + {Role: eschema.User, Content: strings.Repeat("u ", 3000)}, + } + fitted, err := fitExtractorMessages(t.Context(), nil, "test@test", msgs) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(fitted) != 2 { + t.Fatalf("got %d messages, want 2", len(fitted)) + } + if strings.TrimSpace(fitted[1].Content) == "" { + t.Fatal("user turn was emptied") + } +} + +// TestFitExtractorMessages_NoSystemPromptKeepsUserTurn verifies that a +// user-only request (no system prompt configured) is not rejected by the +// system-prompt guard: the guard only applies when a system message was +// actually present, so a valid prompt-only extractor keeps working once the +// model's content_length is resolvable. +func TestFitExtractorMessages_NoSystemPromptKeepsUserTurn(t *testing.T) { + SetExtractorContextLengthOverride(func(_ context.Context, _ string) int { return 2000 }) + t.Cleanup(func() { SetExtractorContextLengthOverride(nil) }) + + msgs := []eschema.Message{ + {Role: eschema.User, Content: strings.Repeat("u ", 3000)}, + } + fitted, err := fitExtractorMessages(t.Context(), nil, "test@test", msgs) + if err != nil { + t.Fatalf("unexpected error for user-only prompt: %v", err) + } + if len(fitted) != 1 || fitted[0].Role != eschema.User { + t.Fatalf("got %d messages, want the single user turn: %+v", len(fitted), fitted) + } + if strings.TrimSpace(fitted[0].Content) == "" { + t.Fatal("user turn was emptied") + } +} + +// TestExtractorComponent_CallRaw_FitsBeforeInvoke verifies the production +// wiring end to end: callRaw resolves the model's context length, trims the +// messages to the budget, and hands the fitted messages to the invoker. +func TestExtractorComponent_CallRaw_FitsBeforeInvoke(t *testing.T) { + SetExtractorContextLengthOverride(func(_ context.Context, _ string) int { return 200 }) + t.Cleanup(func() { SetExtractorContextLengthOverride(nil) }) + + stub := withStubChatInvoker(t, stubResponse{Content: `{"ok": true}`}) + c := &ExtractorComponent{} + + _, err := c.callText(t.Context(), nil, extractorInputs{ + systemPrompt: "extract fields", + prompt: "summarize", + llmID: "test@test", + }, strings.Repeat("chunk text with lots of tokens. ", 500)) + if err != nil { + t.Fatalf("callText: %v", err) + } + + stub.mu.Lock() + req := stub.lastReq + stub.mu.Unlock() + if len(req.Messages) == 0 { + t.Fatal("invoker was not called") + } + if req.Messages[0].Role != eschema.System || strings.TrimSpace(req.Messages[0].Content) == "" { + t.Fatalf("system prompt lost or emptied before invoke: %+v", req.Messages[0]) + } + total := 0 + for _, m := range req.Messages { + total += tokenizer.NumTokensFromString(m.Content) + } + if total > extractorContextFitBudget(200) { + t.Fatalf("sent messages total %d exceed the fitting budget %d", total, extractorContextFitBudget(200)) + } + if !strings.Contains(req.Messages[len(req.Messages)-1].Content, "chunk text") { + t.Fatal("chunk text lost from the user turn") + } +} + +// openExtractorContextTestDB returns an in-memory DB with the tenant and +// tenant-model tables migrated. Tests pass the returned handle explicitly to +// extractorContextLength, defaultChatModelRef, and dao.ResolveModelContentLength, +// so no global DAO state is touched. +func openExtractorContextTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&entity.Tenant{}, &entity.TenantModelProvider{}, &entity.TenantModel{}); err != nil { + t.Fatalf("migrate: %v", err) + } + return db +} + +// seedExtractorContextModel seeds an active OpenAI gpt-4o tenant model +// (catalog content_length 128000) plus its tenant. tenantLLMID, when +// non-empty, pins the tenant's default chat model to the tenant-model UUID; +// otherwise the tenant falls back to the composite llm_id. +func seedExtractorContextModel(t *testing.T, db *gorm.DB, tenantLLMID string) { + t.Helper() + status := "1" + tenant := entity.Tenant{ + ID: "tenant-1", + LLMID: "gpt-4o@openai", + Status: &status, + } + if tenantLLMID != "" { + tenant.TenantLLMID = &tenantLLMID + } + if err := db.Create(&tenant).Error; err != nil { + t.Fatalf("create tenant: %v", err) + } + if err := db.Create(&entity.TenantModelProvider{ + ID: "provider-openai", + ProviderName: "OpenAI", + TenantID: "tenant-1", + }).Error; err != nil { + t.Fatalf("create provider: %v", err) + } + if err := db.Create(&entity.TenantModel{ + ID: "0123456789abcdef0123456789abcdef", + ProviderID: "provider-openai", + InstanceID: "instance-1", + ModelName: "gpt-4o", + ModelType: int(entity.ModelTypeChat), + Status: "active", + }).Error; err != nil { + t.Fatalf("create model: %v", err) + } +} + +// extractorStateCtx returns a context carrying a canvas state with the given +// tenant_id global, as extractorContextLength expects. +func extractorStateCtx(t *testing.T, tenantID string) context.Context { + t.Helper() + state := runtime.NewCanvasState("run-1", "session-1") + state.SetGlobal("tenant_id", tenantID) + return runtime.WithState(t.Context(), state) +} + +// TestExtractorContextLength_TenantModelUUID verifies extractorContextLength +// resolves content_length for a tenant_model UUID through the provider +// catalog. +func TestExtractorContextLength_TenantModelUUID(t *testing.T) { + db := openExtractorContextTestDB(t) + seedExtractorContextModel(t, db, "") + ctx := extractorStateCtx(t, "tenant-1") + + if got := extractorContextLength(ctx, db, "0123456789abcdef0123456789abcdef"); got != 128000 { + t.Fatalf("extractorContextLength(uuid) = %d, want 128000", got) + } +} + +// TestExtractorContextLength_DefaultChatModelPinned verifies the llmID=="" +// fallback resolves the tenant default chat model when it is pinned to a +// tenant_model UUID. +func TestExtractorContextLength_DefaultChatModelPinned(t *testing.T) { + db := openExtractorContextTestDB(t) + seedExtractorContextModel(t, db, "0123456789abcdef0123456789abcdef") + ctx := extractorStateCtx(t, "tenant-1") + + if got := extractorContextLength(ctx, db, ""); got != 128000 { + t.Fatalf("extractorContextLength(default pinned uuid) = %d, want 128000", got) + } +} + +// TestExtractorContextLength_DefaultChatModelComposite verifies the llmID=="" +// fallback resolves the tenant default chat model from the composite +// "model@provider" llm_id when no tenant_model is pinned. +func TestExtractorContextLength_DefaultChatModelComposite(t *testing.T) { + db := openExtractorContextTestDB(t) + seedExtractorContextModel(t, db, "") + ctx := extractorStateCtx(t, "tenant-1") + + if got := extractorContextLength(ctx, db, ""); got != 128000 { + t.Fatalf("extractorContextLength(default composite) = %d, want 128000", got) + } +} + +// TestExtractorContextLength_UnknownModelSkips verifies extractorContextLength +// returns 0 (skip fitting) for an unknown model reference. +func TestExtractorContextLength_UnknownModelSkips(t *testing.T) { + db := openExtractorContextTestDB(t) + seedExtractorContextModel(t, db, "") + ctx := extractorStateCtx(t, "tenant-1") + + if got := extractorContextLength(ctx, db, "no-such-model@no-such-provider"); got != 0 { + t.Fatalf("extractorContextLength(unknown) = %d, want 0", got) + } +} + +// TestExtractorContextFitBudget verifies the fitting budget is 97% of the +// resolved content_length (mirroring the agent's contextFitBudget), leaving +// headroom for tokenizer drift between cl100k and the model's own tokenizer, +// and that a tiny context never collapses to messagefit's <=0 → 8192 default. +func TestExtractorContextFitBudget(t *testing.T) { + if got := extractorContextFitBudget(128000); got != 124160 { + t.Fatalf("extractorContextFitBudget(128000) = %d, want 124160", got) + } + if got := extractorContextFitBudget(1); got != 1 { + t.Fatalf("extractorContextFitBudget(1) = %d, want 1 (clamped to avoid the 8192 Fit default)", got) + } +} + +// TestFitExtractorMessages_RejectsSystemPromptLoss verifies the guard that a +// fitting which empties every system message is rejected instead of sending +// an instruction-less extraction request: the system prompt carries the +// extraction contract, so running with an emptied system prompt would +// silently produce garbage. +func TestFitExtractorMessages_RejectsSystemPromptLoss(t *testing.T) { + SetExtractorContextLengthOverride(func(_ context.Context, _ string) int { return 300 }) + t.Cleanup(func() { SetExtractorContextLengthOverride(nil) }) + + // System dominates (>4x the user) and the user message alone exceeds + // the budget: the proportional trim preserves the user turn and empties + // the system messages. + msgs := []eschema.Message{ + {Role: eschema.System, Content: strings.Repeat("s ", 5000)}, + {Role: eschema.User, Content: strings.Repeat("u ", 400)}, + } + if _, err := fitExtractorMessages(t.Context(), nil, "test@test", msgs); err == nil { + t.Fatal("expected an error when fitting empties the system prompt") + } +} + +// TestExtractorContextLength_NilDBGraceful verifies that resolving the tenant +// default chat model with no database available (nil db and no override) +// degrades to 0 (skip fitting) instead of panicking in defaultChatModelRef. +func TestExtractorContextLength_NilDBGraceful(t *testing.T) { + ctx := extractorStateCtx(t, "tenant-1") + if got := extractorContextLength(ctx, nil, ""); got != 0 { + t.Fatalf("extractorContextLength(nil db, default model) = %d, want 0 (skip fitting)", got) + } +}