From 91fd114783a2153e2d08fc98ff81338546692c5b Mon Sep 17 00:00:00 2001 From: jay77721 <164177721+jay77721@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:32:37 +0800 Subject: [PATCH] fix(dao): honor tenant-configured context window override (D22/D23) (#18171) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make `ResolveModelContentLength` honor the per-model custom **context window length** (`content_length`) — stored in the Python-legacy `tenant_model.extra["max_tokens"]` field, whose semantic meaning is the context window, NOT the generation cap — **before** any provider-catalog read, and remove the parallel service-layer implementation so every consumer shares one resolution path. --- internal/agent/component/llm.go | 10 +- internal/agent/component/llm_test.go | 130 ++++++++- internal/dao/model_context.go | 151 ++++++++-- internal/dao/model_context_test.go | 268 +++++++++++++++--- internal/ingestion/component/extractor.go | 2 +- .../ingestion/component/extractor_test.go | 55 +++- internal/service/model_service.go | 66 +---- internal/service/model_service_test.go | 43 +++ 8 files changed, 595 insertions(+), 130 deletions(-) diff --git a/internal/agent/component/llm.go b/internal/agent/component/llm.go index 77fae6924f..e2ee884918 100644 --- a/internal/agent/component/llm.go +++ b/internal/agent/component/llm.go @@ -314,7 +314,15 @@ func (c *LLMComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[strin // Resolve the model's context window (content_length) for message // fitting. 0 means the model is unknown → fitMessages falls back to // 8192, matching Python's chat_mdl.max_length = model_config.get("max_tokens") or 8192. - contentLength := dao.ResolveModelContentLength(ctx, db, originalModelID, p.Driver, p.ModelID) + // tenantID scopes composite-reference resolution to the tenant's own rows + // so a per-model "max_tokens" override in tenant_model.extra is honored. + tenantID := "" + if state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx); err == nil && state != nil { + if tid, ok := state.Sys["tenant_id"].(string); ok { + tenantID = tid + } + } + contentLength := dao.ResolveModelContentLength(ctx, db, tenantID, originalModelID, p.Driver, p.ModelID) if contentLength <= 0 { // A 0 makes fitMessages fall back to the 8192 default budget, which can // silently discard most of a large-context prompt, so surface the diff --git a/internal/agent/component/llm_test.go b/internal/agent/component/llm_test.go index 5f2a7e6e17..79657751ca 100644 --- a/internal/agent/component/llm_test.go +++ b/internal/agent/component/llm_test.go @@ -262,9 +262,130 @@ func TestLLM_ThinkingFieldRoundTrip(t *testing.T) { } } -// TestLLM_ResolvesTenantModelID guards that custom-added tenant models selected -// in the agent canvas are resolved to their real provider/model name, driver, -// and credentials before the LLM call is dispatched. +// TestLLM_Invoke_CompositeModel_CustomContextOverride verifies the composite +// reference path of the tenant-configured override: a 2000-token extra +// max_tokens on the tenant's gpt-4o row drives trimming even though the +// catalog reports 128k. +func TestLLM_Invoke_CompositeModel_CustomContextOverride(t *testing.T) { + db := setupComponentTestDB(t) + pushComponentDB(t, db) + + if err := db.Create(&entity.TenantModelProvider{ + ID: "provider-comp-1", + TenantID: "tenant-1", + ProviderName: "OpenAI", + }).Error; err != nil { + t.Fatalf("create provider: %v", err) + } + if err := db.Create(&entity.TenantModelInstance{ + ID: "instance-comp-1", + ProviderID: "provider-comp-1", + InstanceName: "default", + APIKey: "test-key", + Status: "active", + }).Error; err != nil { + t.Fatalf("create instance: %v", err) + } + if err := db.Create(&entity.TenantModel{ + ID: "0123456789abcdef0123456789abcdef", + ProviderID: "provider-comp-1", + InstanceID: "instance-comp-1", + ModelName: "gpt-4o", + ModelType: int(entity.ModelTypeChat), + Status: "active", + Extra: `{"max_tokens": 2000}`, + }).Error; err != nil { + t.Fatalf("create model: %v", err) + } + + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "stub"}} + withStubInvoker(t, stub) + + bigPrompt := strings.Repeat("x ", 20000) // ~40k tokens + c := NewLLMComponent(LLMParam{ModelID: "gpt-4o@OpenAI"}) + if _, err := c.Invoke(stateWithTenant("tenant-1"), db, map[string]any{"user_prompt": bigPrompt}); err != nil { + t.Fatalf("Invoke: %v", err) + } + if stub.captured == nil { + t.Fatal("invoker was not called") + } + var userContent string + for _, m := range stub.captured.Messages { + if m.Role == schema.User { + userContent = m.Content + } + } + if userContent == "" { + t.Fatal("no user message captured") + } + if got := tokenizer.NumTokensFromString(userContent); got > 2000 || got < 1000 { + t.Fatalf("user message = %d tokens; want trimmed to the custom 2000-token context window (~1940)", got) + } +} + +// TestLLM_Invoke_UUIDModel_CustomContextOverride verifies end to end that a +// tenant-configured "max_tokens" override in tenant_model.extra wins over the +// provider catalog's content_length: with an override of 2000 and a 40k-token +// prompt, the user message must be trimmed to roughly the override budget, not +// preserved under gpt-4o's 128k catalog window. +func TestLLM_Invoke_UUIDModel_CustomContextOverride(t *testing.T) { + db := setupComponentTestDB(t) + pushComponentDB(t, db) + + if err := db.Create(&entity.TenantModelProvider{ + ID: "provider-uuid-2", + TenantID: "tenant-1", + ProviderName: "OpenAI", + }).Error; err != nil { + t.Fatalf("create provider: %v", err) + } + if err := db.Create(&entity.TenantModelInstance{ + ID: "instance-uuid-2", + ProviderID: "provider-uuid-2", + InstanceName: "default", + APIKey: "test-key", + Status: "active", + }).Error; err != nil { + t.Fatalf("create instance: %v", err) + } + if err := db.Create(&entity.TenantModel{ + ID: "0123456789abcdef0123456789abcdef", + ProviderID: "provider-uuid-2", + InstanceID: "instance-uuid-2", + ModelName: "gpt-4o", + ModelType: int(entity.ModelTypeChat), + Status: "active", + Extra: `{"max_tokens": 2000}`, + }).Error; err != nil { + t.Fatalf("create model: %v", err) + } + + stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "stub"}} + withStubInvoker(t, stub) + + bigPrompt := strings.Repeat("x ", 20000) // ~40k tokens + c := NewLLMComponent(LLMParam{ModelID: "0123456789abcdef0123456789abcdef"}) + if _, err := c.Invoke(stateWithTenant("tenant-1"), db, map[string]any{"user_prompt": bigPrompt}); err != nil { + t.Fatalf("Invoke: %v", err) + } + if stub.captured == nil { + t.Fatal("invoker was not called") + } + var userContent string + for _, m := range stub.captured.Messages { + if m.Role == schema.User { + userContent = m.Content + } + } + if userContent == "" { + t.Fatal("no user message captured") + } + // 97% of the 2000-token override budget; the catalog's 128k must not apply. + if got := tokenizer.NumTokensFromString(userContent); got > 2000 || got < 1000 { + t.Fatalf("user message = %d tokens; want trimmed to the custom 2000-token context window (~1940)", got) + } +} + // TestLLM_Invoke_UUIDModel_ResolvesContentLength verifies the tenant-model // UUID path of content_length resolution end to end: with a real in-memory // DB row for gpt-4o@OpenAI, the fitting budget comes from the catalog's @@ -326,6 +447,9 @@ func TestLLM_Invoke_UUIDModel_ResolvesContentLength(t *testing.T) { } } +// TestLLM_ResolvesTenantModelID guards that custom-added tenant models selected +// in the agent canvas are resolved to their real provider/model name, driver, +// and credentials before the LLM call is dispatched. func TestLLM_ResolvesTenantModelID(t *testing.T) { db := setupComponentTestDB(t) pushComponentDB(t, db) diff --git a/internal/dao/model_context.go b/internal/dao/model_context.go index 616e30f9ca..307eaeb1f0 100644 --- a/internal/dao/model_context.go +++ b/internal/dao/model_context.go @@ -18,45 +18,67 @@ package dao import ( "context" + "encoding/json" + "strconv" "strings" "gorm.io/gorm" + + "ragflow/internal/entity" ) -// ResolveModelContentLength returns the chat model's context window +// ResolveModelContentLength returns the chat model's effective context window // (content_length) in tokens for modelRef — a tenant_model UUID or a // composite "model@provider" / "model@instance@provider" reference — or 0 -// when it cannot be resolved. driver and modelName are an optional provider -// catalog fallback used when modelRef is not a resolvable tenant id (for -// example when no database is available); pass empty strings to skip it. +// when it cannot be resolved. tenantID scopes the composite-reference lookup +// to the tenant's own provider/instance/model rows. driver and modelName are +// an optional provider catalog fallback used when modelRef is not a +// resolvable tenant id (for example when no database is available); pass +// empty strings to skip it. +// +// Resolution order (mirrors Python's model_extra.get("max_tokens") semantics): +// 1. Tenant configuration first: when the tenant_model row is active and +// carries a positive "max_tokens" override in its extra JSON, that custom +// context window wins and no catalog data is read. +// 2. Only with no override, fall back to the provider catalog's +// content_length (via the tenant row, the composite reference parts, or +// the resolved driver + model name). // // This is the shared implementation behind the agent LLM component, the // ingestion Extractor, and service.ModelProviderService — those packages // cannot import each other (import cycles), so the lookup lives here. -func ResolveModelContentLength(ctx context.Context, db *gorm.DB, modelRef, driver, modelName string) int { - // 1. Composite "model@provider" / "model@instance@provider" reference: - // look up the provider catalog directly. A composite reference cannot - // be a tenant-model UUID, so resolve it before touching the database. - if pureName, _, providerName, ok := splitCompositeModelRef(modelRef); ok { +func ResolveModelContentLength(ctx context.Context, db *gorm.DB, tenantID, modelRef, driver, modelName string) int { + if db == nil { + db = DB + } + pureName, instanceName, providerName, composite := splitCompositeModelRef(modelRef) + + // 1. Tenant-configured override: resolve the tenant_model row before any + // catalog read and return the custom context window when present. + obj := lookupTenantModel(ctx, db, tenantID, modelRef, pureName, instanceName, providerName, composite) + if obj != nil && obj.Status == "active" { + if v := modelExtraMaxTokens(obj.Extra); v > 0 { + return v + } + } + + // 2. Catalog fallbacks (only when there is no custom override). + // 2a. Active tenant row → its provider row → catalog content_length. + if obj != nil && obj.Status == "active" { + if provider, err := NewTenantModelProviderDAO().GetByID(ctx, db, obj.ProviderID); err == nil && provider != nil { + if mdl, err := GetModelProviderManager().GetModelByName(provider.ProviderName, obj.ModelName); err == nil && mdl.ContentLength != nil { + return *mdl.ContentLength + } + } + } + // 2b. Composite reference with no tenant row → catalog by reference parts. + if composite { if mdl, err := GetModelProviderManager().GetModelByName(providerName, pureName); err == nil && mdl.ContentLength != nil { return *mdl.ContentLength } } - if db == nil { - db = DB - } - // 2. Tenant model UUID: read content_length from its provider catalog row. - if db != nil && modelRef != "" { - if obj, err := NewTenantModelDAO().GetByID(ctx, db, modelRef); err == nil && obj != nil && obj.Status == "active" { - if provider, err := NewTenantModelProviderDAO().GetByID(ctx, db, obj.ProviderID); err == nil && provider != nil { - if mdl, err := GetModelProviderManager().GetModelByName(provider.ProviderName, obj.ModelName); err == nil && mdl.ContentLength != nil { - return *mdl.ContentLength - } - } - } - } - // 3. Resolved driver + bare model name: fallback when modelRef is a - // tenant id that could not be resolved without a database. + // 2c. Resolved driver + bare model name: fallback when modelRef is a + // tenant id that could not be resolved without a database. if driver != "" && modelName != "" { if mdl, err := GetModelProviderManager().GetModelByName(driver, modelName); err == nil && mdl.ContentLength != nil { return *mdl.ContentLength @@ -65,6 +87,79 @@ func ResolveModelContentLength(ctx context.Context, db *gorm.DB, modelRef, drive return 0 } +// lookupTenantModel resolves the tenant_model row for modelRef — by UUID, or +// for a composite reference through the tenant's provider/instance rows (chat +// model type). Returns nil when the row cannot be resolved. +func lookupTenantModel(ctx context.Context, db *gorm.DB, tenantID, modelRef, pureName, instanceName, providerName string, composite bool) *entity.TenantModel { + if db == nil { + return nil + } + if composite { + if tenantID == "" { + return nil + } + provider, err := NewTenantModelProviderDAO().GetByTenantIDAndProviderName(ctx, db, tenantID, providerName) + if err != nil || provider == nil { + return nil + } + instance, err := NewTenantModelInstanceDAO().GetByProviderIDAndInstanceName(ctx, db, provider.ID, instanceName) + if err != nil || instance == nil || instance.Status != "active" { + return nil + } + obj, err := NewTenantModelDAO().GetByProviderIDAndInstanceIDAndModelTypeAndModelName(ctx, db, provider.ID, instance.ID, int(entity.ModelTypeChat), pureName) + if err != nil || obj == nil { + return nil + } + return obj + } + if modelRef == "" { + return nil + } + obj, err := NewTenantModelDAO().GetByID(ctx, db, modelRef) + if err != nil || obj == nil { + return nil + } + // UUIDs are globally unique and unguessable, and shared/joined-tenant + // models are legitimate (Python's get_model_config_by_id and Go's + // service.GetModelConfigByID both support them), so UUID resolution is + // intentionally NOT scoped to tenantID. The per-model override and the + // catalog fallback apply to whichever tenant holds the UUID. + return obj +} + +// modelExtraMaxTokens returns the per-model "max_tokens" override from the +// tenant_model.extra JSON, or 0 when absent/invalid/non-positive. Mirrors +// Python's model_extra.get("max_tokens") — the custom context-window length a +// user configures for a model. Both JSON numbers and numeric strings are +// accepted (some callers persist extra.max_tokens as a string). +func modelExtraMaxTokens(extra string) int { + if strings.TrimSpace(extra) == "" { + return 0 + } + var m struct { + MaxTokens json.RawMessage `json:"max_tokens"` + } + if err := json.Unmarshal([]byte(extra), &m); err != nil || len(m.MaxTokens) == 0 { + return 0 + } + // Accept JSON numbers (int or float, e.g. 32000 / 32000.0) and numeric + // strings ("32000"); json.Number keeps the raw lexical form. + var num json.Number + if err := json.Unmarshal(m.MaxTokens, &num); err == nil { + if f, err := strconv.ParseFloat(num.String(), 64); err == nil && f > 0 { + return int(f) + } + return 0 + } + var str string + if err := json.Unmarshal(m.MaxTokens, &str); err == nil { + if n, err := strconv.Atoi(strings.TrimSpace(str)); err == nil && n > 0 { + return n + } + } + return 0 +} + // splitCompositeModelRef splits a composite "model@provider" or // "model@instance@provider" reference into its parts. The instance defaults // to "default" for the two-part form. @@ -76,5 +171,13 @@ func splitCompositeModelRef(ref string) (modelName, instanceName, providerName s case 3: return parts[0], parts[1], parts[2], true } + if len(parts) > 3 { + // 4+ segments: any '@' embedded in the leftmost modelName component + // must be preserved (e.g. LM Studio chat models + // `name@q8_0@lmstudio@LM-Studio`). Rejoin the leading fields into the + // model name, keeping instance and provider anchored on the right. + n := len(parts) + return strings.Join(parts[:n-2], "@"), parts[n-2], parts[n-1], true + } return "", "", "", false } diff --git a/internal/dao/model_context_test.go b/internal/dao/model_context_test.go index 54931c1460..aa397f8765 100644 --- a/internal/dao/model_context_test.go +++ b/internal/dao/model_context_test.go @@ -28,7 +28,7 @@ import ( // TestResolveModelContentLength_CompositeReference resolves content_length // for a composite "model@provider" reference from the provider catalog. func TestResolveModelContentLength_CompositeReference(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "gpt-4o@openai", "", ""); got != 128000 { + if got := ResolveModelContentLength(t.Context(), nil, "", "gpt-4o@openai", "", ""); got != 128000 { t.Fatalf("ResolveModelContentLength(gpt-4o@openai) = %d, want 128000", got) } } @@ -38,43 +38,44 @@ func TestResolveModelContentLength_CompositeReference(t *testing.T) { // instance-bearing form used by tenant model instances — from the provider // catalog. func TestResolveModelContentLength_CompositeReferenceThreePart(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "gpt-4o@default@openai", "", ""); got != 128000 { + if got := ResolveModelContentLength(t.Context(), nil, "", "gpt-4o@default@openai", "", ""); got != 128000 { t.Fatalf("ResolveModelContentLength(gpt-4o@default@openai) = %d, want 128000", got) } } -// TestResolveModelContentLength_TooManyParts documents that a reference with -// more than two "@" separators is not treated as a composite ref: it falls -// through to the driver+modelName fallback, and to 0 when no fallback is -// supplied. A known catalog model with an extra separator is used for the -// no-fallback assertion so a parser regression (accepting excessive -// separators) fails loudly instead of resolving to 0 by coincidence. -func TestResolveModelContentLength_TooManyParts(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "a@b@c@d", "openai", "gpt-4o"); got != 128000 { +// TestResolveModelContentLength_MultiSegmentComposite documents that a +// reference with more than three "@" segments is still a composite ref: the +// leading segments are rejoined into the model name (preserving embedded '@', +// e.g. LM Studio chat models `name@q8_0@lmstudio@LM-Studio`), with instance +// and provider anchored on the right. When the catalog lookup fails it falls +// through to the driver+modelName fallback, and to 0 when no fallback exists. +func TestResolveModelContentLength_MultiSegmentComposite(t *testing.T) { + if got := ResolveModelContentLength(t.Context(), nil, "", "a@b@c@d", "openai", "gpt-4o"); got != 128000 { t.Fatalf("ResolveModelContentLength(a@b@c@d, openai, gpt-4o) = %d, want 128000 (driver fallback)", got) } - if got := ResolveModelContentLength(t.Context(), nil, "gpt-4o@default@openai@extra", "", ""); got != 0 { + if got := ResolveModelContentLength(t.Context(), nil, "", "gpt-4o@default@openai@extra", "", ""); got != 0 { t.Fatalf("ResolveModelContentLength(gpt-4o@default@openai@extra) = %d, want 0", got) } - if _, _, _, ok := splitCompositeModelRef("a@b@c@d"); ok { - t.Fatal("splitCompositeModelRef(a@b@c@d) accepted an excessive-separator reference") + modelName, instanceName, providerName, ok := splitCompositeModelRef("a@b@c@d") + if !ok || modelName != "a@b" || instanceName != "c" || providerName != "d" { + t.Fatalf("splitCompositeModelRef(a@b@c@d) = %q/%q/%q/%v, want a@b/c/d/true", modelName, instanceName, providerName, ok) } } // TestResolveModelContentLength_DriverModelFallback resolves content_length // from the resolved driver + bare model name, which needs no database. func TestResolveModelContentLength_DriverModelFallback(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "", "openai", "gpt-4o"); got != 128000 { + if got := ResolveModelContentLength(t.Context(), nil, "", "", "openai", "gpt-4o"); got != 128000 { t.Fatalf("ResolveModelContentLength(openai/gpt-4o) = %d, want 128000", got) } } // TestResolveModelContentLength_Unknown returns 0 for unknown references. func TestResolveModelContentLength_Unknown(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "no-such-model@no-such-provider", "", ""); got != 0 { + if got := ResolveModelContentLength(t.Context(), nil, "", "no-such-model@no-such-provider", "", ""); got != 0 { t.Fatalf("unknown model = %d, want 0", got) } - if got := ResolveModelContentLength(t.Context(), nil, "", "", ""); got != 0 { + if got := ResolveModelContentLength(t.Context(), nil, "", "", "", ""); got != 0 { t.Fatalf("empty reference = %d, want 0", got) } } @@ -86,58 +87,233 @@ func TestResolveModelContentLength_TenantModelUUID(t *testing.T) { pushDB(t, db) ctx := t.Context() - 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) - } + seedOpenAIChatModel(t, db, "") - if got := ResolveModelContentLength(ctx, db, "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { + if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { t.Fatalf("ResolveModelContentLength(uuid) = %d, want 128000", got) } } +// TestResolveModelContentLength_ExtraOverrideUUID verifies that a custom +// "max_tokens" context window in the tenant_model.extra wins over the +// provider catalog's content_length for a UUID reference. +func TestResolveModelContentLength_ExtraOverrideUUID(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 32000 { + t.Fatalf("ResolveModelContentLength(uuid+extra) = %d, want 32000 (custom override wins over catalog 128000)", got) + } +} + +// TestResolveModelContentLength_ExtraOverrideComposite verifies the custom +// context window override for a composite reference resolved through the +// tenant's provider/instance/model rows. +func TestResolveModelContentLength_ExtraOverrideComposite(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 32000 { + t.Fatalf("ResolveModelContentLength(composite+extra) = %d, want 32000", got) + } +} + +// TestResolveModelContentLength_CustomModelExtraComposite is the core +// custom-model scenario: a model name that is NOT in the provider catalog but +// carries a tenant-configured "max_tokens" override must resolve to that +// override (the catalog cannot provide a value). +func TestResolveModelContentLength_CustomModelExtraComposite(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedCustomChatModel(t, db, "my-local-model", `{"max_tokens": 32000}`) + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "my-local-model@OpenAI", "", ""); got != 32000 { + t.Fatalf("ResolveModelContentLength(custom+extra) = %d, want 32000", got) + } + // Without the override the custom model is unknown to the catalog → 0. + if err := db.Model(&entity.TenantModel{}). + Where("id = ?", "0123456789abcdef0123456789abcdef"). + Update("extra", "").Error; err != nil { + t.Fatalf("clear custom model extra: %v", err) + } + if got := ResolveModelContentLength(ctx, db, "tenant-1", "my-local-model@OpenAI", "", ""); got != 0 { + t.Fatalf("custom model without extra = %d, want 0", got) + } +} + +// TestResolveModelContentLength_InactiveInstanceCompositeFallsBackToCatalog +// verifies that a composite reference whose instance is inactive does not +// apply the override or the tenant-row catalog read: it falls back to the +// catalog by reference parts (gpt-4o → 128000). +func TestResolveModelContentLength_InactiveInstanceCompositeFallsBackToCatalog(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) + if err := db.Model(&entity.TenantModelInstance{}). + Where("id = ?", "instance-1"). + Update("status", "inactive").Error; err != nil { + t.Fatalf("set instance inactive: %v", err) + } + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { + t.Fatalf("inactive-instance composite = %d, want catalog 128000", got) + } +} + +// TestResolveModelContentLength_InactiveModelCompositeFallsBackToCatalog +// verifies that a composite reference whose model row is inactive does not +// apply the override or the tenant-row catalog read; it falls back to the +// catalog by reference parts (gpt-4o → 128000). +func TestResolveModelContentLength_InactiveModelCompositeFallsBackToCatalog(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) + if err := db.Model(&entity.TenantModel{}). + Where("id = ?", "0123456789abcdef0123456789abcdef"). + Update("status", "inactive").Error; err != nil { + t.Fatalf("set model inactive: %v", err) + } + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { + t.Fatalf("inactive-model composite = %d, want catalog 128000", got) + } +} + +// TestResolveModelContentLength_ExtraOverrideStringForm verifies that a +// max_tokens override persisted as a JSON string is still honored (some +// callers persist extra.max_tokens as a string). +func TestResolveModelContentLength_ExtraOverrideStringForm(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedOpenAIChatModel(t, db, `{"max_tokens": "32000"}`) + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 32000 { + t.Fatalf("ResolveModelContentLength(string-form extra) = %d, want 32000", got) + } +} + +// TestResolveModelContentLength_NonNumericStringExtraFallsBackToCatalog +// verifies that a non-numeric max_tokens string is treated as "no override". +func TestResolveModelContentLength_NonNumericStringExtraFallsBackToCatalog(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedOpenAIChatModel(t, db, `{"max_tokens": "huge"}`) + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { + t.Fatalf("ResolveModelContentLength(non-numeric string extra) = %d, want catalog 128000", got) + } +} + +// TestResolveModelContentLength_InvalidExtraFallsBackToCatalog verifies that +// an unparsable extra JSON is treated as "no override" and the catalog +// content_length is used. +func TestResolveModelContentLength_InvalidExtraFallsBackToCatalog(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedOpenAIChatModel(t, db, `{bad json`) + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { + t.Fatalf("ResolveModelContentLength(invalid extra) = %d, want catalog 128000", got) + } +} + +// TestResolveModelContentLength_ZeroExtraFallsBackToCatalog verifies that a +// non-positive max_tokens override is ignored in favor of the catalog. +func TestResolveModelContentLength_ZeroExtraFallsBackToCatalog(t *testing.T) { + db := openModelContextTestDB(t) + pushDB(t, db) + ctx := t.Context() + + seedOpenAIChatModel(t, db, `{"max_tokens": 0}`) + + if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { + t.Fatalf("ResolveModelContentLength(zero extra) = %d, want catalog 128000", got) + } +} + // TestResolveModelContentLength_InactiveTenantModel falls through to the -// composite/catalog paths when the tenant model row is not active. +// composite/catalog paths when the tenant model row is not active — even when +// it carries a max_tokens override. func TestResolveModelContentLength_InactiveTenantModel(t *testing.T) { db := openModelContextTestDB(t) pushDB(t, db) ctx := t.Context() + seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) + // Flip the row to inactive. + if err := db.Model(&entity.TenantModel{}). + Where("id = ?", "0123456789abcdef0123456789abcdef"). + Update("status", "inactive").Error; err != nil { + t.Fatalf("set inactive: %v", err) + } + + // The UUID is not a composite ref, so an inactive row yields 0. + if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 0 { + t.Fatalf("inactive tenant model = %d, want 0", got) + } +} + +// seedOpenAIChatModel seeds an active OpenAI gpt-4o tenant model (catalog +// content_length 128000) plus its provider and default instance. extra is the +// tenant_model.extra JSON ("" for none). +func seedOpenAIChatModel(t *testing.T, db *gorm.DB, extra string) { + t.Helper() + seedChatModel(t, db, "provider-openai", "OpenAI", "gpt-4o", extra) +} + +// seedCustomChatModel seeds an active tenant model whose name is NOT in the +// provider catalog (custom/local model scenario). +func seedCustomChatModel(t *testing.T, db *gorm.DB, modelName, extra string) { + t.Helper() + seedChatModel(t, db, "provider-openai", "OpenAI", modelName, extra) +} + +func seedChatModel(t *testing.T, db *gorm.DB, providerID, providerName, modelName, extra string) { + t.Helper() if err := db.Create(&entity.TenantModelProvider{ - ID: "provider-openai", - ProviderName: "OpenAI", + ID: providerID, + ProviderName: providerName, TenantID: "tenant-1", }).Error; err != nil { t.Fatalf("create provider: %v", err) } + if err := db.Create(&entity.TenantModelInstance{ + ID: "instance-1", + ProviderID: providerID, + InstanceName: "default", + Status: "active", + }).Error; err != nil { + t.Fatalf("create instance: %v", err) + } if err := db.Create(&entity.TenantModel{ ID: "0123456789abcdef0123456789abcdef", - ProviderID: "provider-openai", + ProviderID: providerID, InstanceID: "instance-1", - ModelName: "gpt-4o", + ModelName: modelName, ModelType: int(entity.ModelTypeChat), - Status: "inactive", + Status: "active", + Extra: extra, }).Error; err != nil { t.Fatalf("create model: %v", err) } - - // The UUID is not a composite ref, so an inactive row yields 0. - if got := ResolveModelContentLength(ctx, db, "0123456789abcdef0123456789abcdef", "", ""); got != 0 { - t.Fatalf("inactive tenant model = %d, want 0", got) - } } func openModelContextTestDB(t *testing.T) *gorm.DB { @@ -146,7 +322,7 @@ func openModelContextTestDB(t *testing.T) *gorm.DB { if err != nil { t.Fatalf("open sqlite: %v", err) } - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatalf("migrate: %v", err) } return db diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index f9f1295644..9d72d50e24 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -1416,7 +1416,7 @@ func extractorContextLength(ctx context.Context, db *gorm.DB, llmID string) int if llmID == "" { return 0 } - return dao.ResolveModelContentLength(ctx, db, llmID, "", "") + return dao.ResolveModelContentLength(ctx, db, tid, llmID, "", "") } // defaultChatModelRef returns the tenant's default chat model reference — diff --git a/internal/ingestion/component/extractor_test.go b/internal/ingestion/component/extractor_test.go index 0e882b23df..80450e703d 100644 --- a/internal/ingestion/component/extractor_test.go +++ b/internal/ingestion/component/extractor_test.go @@ -1901,6 +1901,59 @@ func TestExtractorComponent_CallRaw_FitsBeforeInvoke(t *testing.T) { } } +// TestExtractorComponent_CallRaw_CustomContextOverride verifies the extractor +// wiring honors the tenant-configured override end to end: with a 2000-token +// extra max_tokens on the gpt-4o row, the invoker receives messages fitted to +// ~1940 tokens instead of the catalog's 128k. +func TestExtractorComponent_CallRaw_CustomContextOverride(t *testing.T) { + db := openExtractorContextTestDB(t) + seedExtractorContextModel(t, db, "") + // Add the instance row the composite resolution path needs, then pin the + // tenant-configured context override on the model. + if err := db.Create(&entity.TenantModelInstance{ + ID: "instance-1", + ProviderID: "provider-openai", + InstanceName: "default", + Status: "active", + }).Error; err != nil { + t.Fatalf("create instance: %v", err) + } + if err := db.Model(&entity.TenantModel{}). + Where("id = ?", "0123456789abcdef0123456789abcdef"). + Update("extra", `{"max_tokens": 2000}`).Error; err != nil { + t.Fatalf("set model extra: %v", err) + } + ctx := extractorStateCtx(t, "tenant-1") + + stub := withStubChatInvoker(t, stubResponse{Content: `{"ok": true}`}) + c := &ExtractorComponent{} + _, err := c.callText(ctx, db, extractorInputs{ + systemPrompt: "extract fields", + prompt: "summarize", + llmID: "gpt-4o@OpenAI", + }, 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: %+v", req.Messages[0]) + } + total := 0 + for _, m := range req.Messages { + total += tokenizer.NumTokensFromString(m.Content) + } + if total > 2000 { + t.Fatalf("sent messages total %d exceed the custom 2000-token context window", total) + } +} + // 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, @@ -1911,7 +1964,7 @@ func openExtractorContextTestDB(t *testing.T) *gorm.DB { if err != nil { t.Fatalf("open sqlite: %v", err) } - if err := db.AutoMigrate(&entity.Tenant{}, &entity.TenantModelProvider{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.Tenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatalf("migrate: %v", err) } return db diff --git a/internal/service/model_service.go b/internal/service/model_service.go index adb27d2eaf..e7cc0afb63 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -3594,64 +3594,22 @@ func (m *ModelProviderService) ResolveModelConfig(ctx context.Context, tenantID return m.GetModelConfigFromProviderInstance(ctx, tenantID, modelType, modelRef) } -// ResolveModelContextLength returns the chat model's context window -// (content_length) in tokens, or 0 when unknown. After the all_models.json -// migration (PR #17839) content_length is the total context window and -// max_output is the generation cap; the knowledge_compiler prompt-budget logic -// needs the context window, not the output cap. modelRef accepts either a -// tenant model UUID or a "model@instance@provider" composite name. +// ResolveModelContextLength returns the chat model's effective context window +// (content_length) in tokens, or 0 when unknown. content_length is the total +// context window and max_output is the generation cap; the +// knowledge_compiler prompt-budget logic needs the context window, not the +// output cap. modelRef accepts either a tenant model UUID or a +// "model@instance@provider" composite name. +// +// The resolution is delegated to dao.ResolveModelContentLength so every +// consumer shares one path: a tenant-configured "max_tokens" override in the +// tenant_model.extra wins, otherwise the provider catalog's content_length is +// used (D22/D23). func (m *ModelProviderService) ResolveModelContextLength(ctx context.Context, tenantID string, modelRef string) (int, error) { if strings.TrimSpace(modelRef) == "" { return 0, fmt.Errorf("model ref is required") } - if modelObj, err := m.modelDAO.GetByID(ctx, dao.DB, modelRef); err == nil { - return m.modelContextLengthByID(ctx, modelObj) - } else if !errors.Is(err, gorm.ErrRecordNotFound) { - return 0, err - } - pureName, _, providerName, err := parseModelName(modelRef) - if err != nil { - return 0, err - } - return m.modelContextLengthByName(providerName, pureName) -} - -// modelContextLengthByID reads content_length from the factory catalog for a -// tenant model row (by id). -func (m *ModelProviderService) modelContextLengthByID(ctx context.Context, modelObj *entity.TenantModel) (int, error) { - if modelObj.Status != "active" { - return 0, fmt.Errorf("tenant model id=%s is disabled", modelObj.ID) - } - provider, err := m.modelProviderDAO.GetByID(ctx, dao.DB, modelObj.ProviderID) - if err != nil { - return 0, err - } - if provider == nil { - return 0, fmt.Errorf("provider id=%s not found for model id=%s", modelObj.ProviderID, modelObj.ID) - } - if mi, _ := dao.GetModelProviderManager().GetModelByName(provider.ProviderName, modelObj.ModelName); mi != nil && mi.ContentLength != nil { - return *mi.ContentLength, nil - } - return 0, nil -} - -// modelContextLengthByName reads content_length from the factory catalog for a -// "model@provider" style reference. It is best-effort: an unknown provider or -// model returns 0 (caller falls back to a default context length). -func (m *ModelProviderService) modelContextLengthByName(providerName, pureName string) (int, error) { - targetProvider := dao.GetModelProviderManager().FindProvider(providerName) - if targetProvider == nil { - return 0, fmt.Errorf("model provider config not found: %s", providerName) - } - for i := range targetProvider.Models { - if strings.EqualFold(targetProvider.Models[i].Name, pureName) { - if targetProvider.Models[i].ContentLength != nil { - return *targetProvider.Models[i].ContentLength, nil - } - return 0, nil - } - } - return 0, nil + return dao.ResolveModelContentLength(ctx, dao.DB, tenantID, modelRef, "", ""), nil } func (m *ModelProviderService) ResolveModelID(ctx context.Context, tenantID string, modelType entity.ModelType, modelName string) (string, error) { diff --git a/internal/service/model_service_test.go b/internal/service/model_service_test.go index 86cc46cc3d..c84081cfc9 100644 --- a/internal/service/model_service_test.go +++ b/internal/service/model_service_test.go @@ -373,6 +373,49 @@ func TestModelProviderServiceResolveModelContextLengthUnknownModel(t *testing.T) } } +// TestModelProviderServiceResolveModelContextLengthOverride verifies that the +// tenant-configured "max_tokens" override in tenant_model.extra wins over the +// catalog content_length through the service delegation. UUID resolution is +// unscoped (globally unique); the composite path needs the real tenant id to +// locate the tenant's provider/instance/model rows. +func TestModelProviderServiceResolveModelContextLengthOverride(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + useModelProviderServiceTestDB(t, db) + activeStatus := "1" + rows := []interface{}{ + &entity.UserTenant{ID: "user-tenant-cl", UserID: "user-1", TenantID: "tenant-cl", Role: "owner", InvitedBy: "user-1", Status: &activeStatus}, + &entity.TenantModelProvider{ID: "provider-anthropic", TenantID: "tenant-cl", ProviderName: "Anthropic"}, + &entity.TenantModelInstance{ID: "instance-anthropic", ProviderID: "provider-anthropic", InstanceName: "default", APIKey: "sk-anthropic", Status: "active", Extra: "{}"}, + &entity.TenantModel{ID: "model-claude", ProviderID: "provider-anthropic", InstanceID: "instance-anthropic", ModelName: "claude-opus-4-8", ModelType: int(entity.ModelTypeChat), Status: "active", Extra: `{"max_tokens": 4096}`}, + } + for _, row := range rows { + if err := db.Create(row).Error; err != nil { + t.Fatalf("failed to seed %T: %v", row, err) + } + } + + svc := NewModelProviderService() + ctx := t.Context() + + // UUID path: the 4096 override wins over catalog content_length 1000000. + got, err := svc.ResolveModelContextLength(ctx, "user-1", "model-claude") + if err != nil { + t.Fatalf("ResolveModelContextLength(override uuid) error = %v", err) + } + if got != 4096 { + t.Fatalf("uuid override content_length = %d, want 4096 (custom override, not catalog 1000000)", got) + } + + // Composite path with the real tenant id honors the same override. + got2, err := svc.ResolveModelContextLength(ctx, "tenant-cl", "claude-opus-4-8@default@Anthropic") + if err != nil { + t.Fatalf("ResolveModelContextLength(override composite) error = %v", err) + } + if got2 != 4096 { + t.Fatalf("composite override content_length = %d, want 4096", got2) + } +} + func TestModelProviderServiceAlterModelRejectsInvalidStatus(t *testing.T) { ctx := t.Context() code, err := NewModelProviderService().AlterModel(ctx, "OpenAI", "default", "", "user-1", "model-1", map[string]interface{}{"status": "disabled"})