ragflow/internal/dao/model_context_test.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package dao

import (
	"testing"

	"github.com/glebarez/sqlite"
	"gorm.io/gorm"

	"ragflow/internal/entity"
)

// 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 {
		t.Fatalf("ResolveModelContentLength(gpt-4o@openai) = %d, want 128000", got)
	}
}

// TestResolveModelContentLength_CompositeReferenceThreePart resolves
// content_length for a composite "model@instance@provider" reference — the
// 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 {
		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 {
		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 {
		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")
	}
}

// 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 {
		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 {
		t.Fatalf("unknown model = %d, want 0", got)
	}
	if got := ResolveModelContentLength(t.Context(), nil, "", "", ""); got != 0 {
		t.Fatalf("empty reference = %d, want 0", got)
	}
}

// TestResolveModelContentLength_TenantModelUUID resolves content_length for a
// tenant_model UUID row through the provider catalog.
func TestResolveModelContentLength_TenantModelUUID(t *testing.T) {
	db := openModelContextTestDB(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)
	}

	if got := ResolveModelContentLength(ctx, db, "0123456789abcdef0123456789abcdef", "", ""); got != 128000 {
		t.Fatalf("ResolveModelContentLength(uuid) = %d, want 128000", got)
	}
}

// TestResolveModelContentLength_InactiveTenantModel falls through to the
// composite/catalog paths when the tenant model row is not active.
func TestResolveModelContentLength_InactiveTenantModel(t *testing.T) {
	db := openModelContextTestDB(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:     "inactive",
	}).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 {
	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.TenantModelProvider{}, &entity.TenantModel{}); err != nil {
		t.Fatalf("migrate: %v", err)
	}
	return db
}