mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-14 12:47:14 +08:00
feat(component): add shared message-fit package (#18091)
The agent LLM component and the ingestion Extractor component both need to trim prompts to the model's context window before calling the provider. Each previously did (or would do) this with its own copy of the logic. This PR adds the shared primitive; follow-up PRs wire it into the agent LLM component (#18092) and the ingestion Extractor/tagger (#18095).
This commit is contained in:
jay77721
183
internal/component/messagefit/messagefit.go
Normal file
183
internal/component/messagefit/messagefit.go
Normal file
@@ -0,0 +1,183 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (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
|
||||
//
|
||||
|
||||
// Package messagefit trims a message list so its total token count fits
|
||||
// within a budget. It mirrors Python's rag/prompts/generator.py:message_fit_in.
|
||||
//
|
||||
// The package is shared by the agent canvas LLM component and the ingestion
|
||||
// Extractor component. Callers convert their message type to []Message, call
|
||||
// Fit, and send the returned kept messages (at keptIdx into the input).
|
||||
package messagefit
|
||||
|
||||
import (
|
||||
"slices"
|
||||
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
// Message is the minimal representation the fitter needs. Both the agent's
|
||||
// schema.Message and ingestion's eschema.Message convert to this.
|
||||
type Message struct {
|
||||
// Role is "system", "user", "assistant", etc. Only "system" receives
|
||||
// special treatment during fitting.
|
||||
Role string
|
||||
// Content is the text content that may be truncated.
|
||||
Content string
|
||||
}
|
||||
|
||||
// Fit trims msgs so the kept messages fit within budget. budget is the
|
||||
// caller-chosen token ceiling for the whole conversation — the agent LLM and
|
||||
// ingestion Extractor both pass the chat model's context window
|
||||
// (content_length), not the generation cap (max_output).
|
||||
//
|
||||
// It returns the kept messages in original order (trimmed when necessary),
|
||||
// their original indices into msgs, and the kept messages' total token count.
|
||||
// msgs itself is never modified, and dropped entries are simply absent from
|
||||
// kept/keptIdx — no empty-content sentinel is used. A message that is kept but
|
||||
// trimmed to empty (e.g. the system share collapses to 0 when the last message
|
||||
// alone fills the budget) is still reported as kept, mirroring Python.
|
||||
//
|
||||
// Strategy (mirrors Python's message_fit_in, with two deliberate tweaks:
|
||||
// an exact budget match counts as fitting, and the system share is spread
|
||||
// across every retained system message instead of only the first):
|
||||
// 1. If everything fits, return as-is.
|
||||
// 2. Keep all system messages + the last non-system message, drop the
|
||||
// rest; if that fits, return.
|
||||
// 3. If still over, trim proportionally:
|
||||
// - System dominates (>80% of tokens) → preserve the last message,
|
||||
// give the remaining budget to the system messages.
|
||||
// - Otherwise → preserve the system messages, give the remaining
|
||||
// budget to the last.
|
||||
// - Single message → trim to budget directly.
|
||||
//
|
||||
// budget <= 0 is treated as 8192 (Python's default).
|
||||
func Fit(msgs []Message, budget int) (kept []Message, keptIdx []int, count int) {
|
||||
if budget <= 0 {
|
||||
budget = 8192
|
||||
}
|
||||
if len(msgs) == 0 {
|
||||
return nil, nil, 0
|
||||
}
|
||||
|
||||
// Step 1: everything fits (an exact budget match counts as fitting).
|
||||
if total := countTokens(msgs); total <= budget {
|
||||
kept = slices.Clone(msgs)
|
||||
keptIdx = make([]int, len(msgs))
|
||||
for i := range keptIdx {
|
||||
keptIdx[i] = i
|
||||
}
|
||||
return kept, keptIdx, total
|
||||
}
|
||||
|
||||
// Step 2: keep all system + last non-system.
|
||||
kept = make([]Message, 0, len(msgs))
|
||||
keptIdx = make([]int, 0, len(msgs))
|
||||
lastNonSystem := -1
|
||||
for i := range msgs {
|
||||
if msgs[i].Role == "system" {
|
||||
kept = append(kept, msgs[i])
|
||||
keptIdx = append(keptIdx, i)
|
||||
} else {
|
||||
lastNonSystem = i
|
||||
}
|
||||
}
|
||||
if lastNonSystem >= 0 {
|
||||
kept = append(kept, msgs[lastNonSystem])
|
||||
keptIdx = append(keptIdx, lastNonSystem)
|
||||
}
|
||||
if len(kept) == 0 {
|
||||
return nil, nil, 0
|
||||
}
|
||||
if total := countTokens(kept); total <= budget {
|
||||
return kept, keptIdx, total
|
||||
}
|
||||
|
||||
// Step 3: trim proportionally.
|
||||
if len(kept) == 1 {
|
||||
kept[0].Content = tokenizer.TrimContentToTokenLimit(kept[0].Content, budget)
|
||||
return kept, keptIdx, countTokens(kept)
|
||||
}
|
||||
|
||||
// Only system messages were retained (no non-system message): spread the
|
||||
// whole budget across every retained system message.
|
||||
if lastNonSystem < 0 {
|
||||
trimSystems(kept, budget)
|
||||
return kept, keptIdx, countTokens(kept)
|
||||
}
|
||||
|
||||
// kept[:len(kept)-1] are the retained system messages; the last entry
|
||||
// is the final non-system message.
|
||||
sys := kept[:len(kept)-1]
|
||||
last := &kept[len(kept)-1]
|
||||
ll := 0
|
||||
for i := range sys {
|
||||
ll += tokenizer.NumTokensFromString(sys[i].Content)
|
||||
}
|
||||
ll2 := tokenizer.NumTokensFromString(last.Content)
|
||||
total := ll + ll2
|
||||
if total <= 0 {
|
||||
return kept, keptIdx, 0
|
||||
}
|
||||
|
||||
if float64(ll)/float64(total) > 0.8 {
|
||||
// System dominates: preserve the last message and give the
|
||||
// remaining budget to the system messages.
|
||||
preserved := min(ll2, budget)
|
||||
last.Content = tokenizer.TrimContentToTokenLimit(last.Content, preserved)
|
||||
trimSystems(sys, max(0, budget-preserved))
|
||||
} else {
|
||||
preserved := min(ll, budget)
|
||||
trimSystems(sys, preserved)
|
||||
last.Content = tokenizer.TrimContentToTokenLimit(last.Content, max(0, budget-preserved))
|
||||
}
|
||||
return kept, keptIdx, countTokens(kept)
|
||||
}
|
||||
|
||||
// trimSystems trims each system message so their combined token count fits
|
||||
// within budget. The budget is allocated in proportion to each message's
|
||||
// original token count, with the last message taking any remainder so the
|
||||
// total never exceeds budget.
|
||||
func trimSystems(sys []Message, budget int) {
|
||||
if len(sys) == 0 {
|
||||
return
|
||||
}
|
||||
if budget <= 0 {
|
||||
for i := range sys {
|
||||
sys[i].Content = ""
|
||||
}
|
||||
return
|
||||
}
|
||||
total := 0
|
||||
for i := range sys {
|
||||
total += tokenizer.NumTokensFromString(sys[i].Content)
|
||||
}
|
||||
if total <= 0 {
|
||||
return
|
||||
}
|
||||
remaining := budget
|
||||
for i := range sys {
|
||||
limit := remaining
|
||||
if i < len(sys)-1 {
|
||||
tokens := tokenizer.NumTokensFromString(sys[i].Content)
|
||||
limit = int(float64(budget) * float64(tokens) / float64(total))
|
||||
if limit > remaining {
|
||||
limit = remaining
|
||||
}
|
||||
}
|
||||
sys[i].Content = tokenizer.TrimContentToTokenLimit(sys[i].Content, limit)
|
||||
remaining -= tokenizer.NumTokensFromString(sys[i].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func countTokens(msgs []Message) int {
|
||||
total := 0
|
||||
for i := range msgs {
|
||||
total += tokenizer.NumTokensFromString(msgs[i].Content)
|
||||
}
|
||||
return total
|
||||
}
|
||||
297
internal/component/messagefit/messagefit_test.go
Normal file
297
internal/component/messagefit/messagefit_test.go
Normal file
@@ -0,0 +1,297 @@
|
||||
package messagefit
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
func TestFit_AllFits(t *testing.T) {
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: "hello"},
|
||||
{Role: "user", Content: "world"},
|
||||
}
|
||||
kept, keptIdx, count := Fit(msgs, 100000)
|
||||
if count == 0 {
|
||||
t.Errorf("Fit returned count 0, want > 0")
|
||||
}
|
||||
if len(kept) != 2 || !slices.Equal(keptIdx, []int{0, 1}) {
|
||||
t.Fatalf("got kept=%+v keptIdx=%v, want both messages", kept, keptIdx)
|
||||
}
|
||||
if kept[0].Content != "hello" || kept[1].Content != "world" {
|
||||
t.Errorf("messages modified when they fit: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFit_DoesNotMutateInput verifies the non-mutating contract: Fit never
|
||||
// rewrites msgs, so a caller cannot accidentally send emptied entries.
|
||||
func TestFit_DoesNotMutateInput(t *testing.T) {
|
||||
orig := []Message{
|
||||
{Role: "system", Content: strings.Repeat("s ", 500)},
|
||||
{Role: "user", Content: strings.Repeat("u ", 500)},
|
||||
{Role: "user", Content: "last"},
|
||||
}
|
||||
msgs := slices.Clone(orig)
|
||||
Fit(msgs, 100)
|
||||
if !slices.Equal(msgs, orig) {
|
||||
t.Fatalf("Fit mutated its input:\n got %+v\nwant %+v", msgs, orig)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_Step2_DropsMiddle(t *testing.T) {
|
||||
// system + last user fit within the budget, but all three together do
|
||||
// not, so Step 2 drops the middle user and keeps system + last intact.
|
||||
sysContent := strings.Repeat("x ", 200)
|
||||
middle := "middle"
|
||||
last := "last"
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: sysContent},
|
||||
{Role: "user", Content: middle},
|
||||
{Role: "user", Content: last},
|
||||
}
|
||||
budget := tokenizer.NumTokensFromString(sysContent) + tokenizer.NumTokensFromString(last)
|
||||
|
||||
kept, keptIdx, count := Fit(msgs, budget)
|
||||
if count == 0 {
|
||||
t.Fatalf("Fit returned count 0, want > 0")
|
||||
}
|
||||
if !slices.Equal(keptIdx, []int{0, 2}) {
|
||||
t.Fatalf("keptIdx = %v, want [0 2] (middle dropped)", keptIdx)
|
||||
}
|
||||
if kept[0].Content != sysContent || kept[1].Content != last {
|
||||
t.Errorf("retained messages modified: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_ExactBudget(t *testing.T) {
|
||||
// A total exactly equal to the budget counts as fitting.
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: "abc"},
|
||||
{Role: "user", Content: "def"},
|
||||
}
|
||||
budget := tokenizer.NumTokensFromString("abc") + tokenizer.NumTokensFromString("def")
|
||||
|
||||
kept, keptIdx, _ := Fit(msgs, budget)
|
||||
if len(kept) != 2 || !slices.Equal(keptIdx, []int{0, 1}) {
|
||||
t.Fatalf("got kept=%+v keptIdx=%v, want both messages kept", kept, keptIdx)
|
||||
}
|
||||
if kept[0].Content != "abc" || kept[1].Content != "def" {
|
||||
t.Errorf("messages modified at exact budget: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFit_NoSystem_KeepsOnlyLast locks the Python-parity behavior: without a
|
||||
// system message, Step 2 keeps only the last non-system message.
|
||||
func TestFit_NoSystem_KeepsOnlyLast(t *testing.T) {
|
||||
msgs := []Message{
|
||||
{Role: "user", Content: strings.Repeat("a ", 300)},
|
||||
{Role: "assistant", Content: strings.Repeat("b ", 300)},
|
||||
{Role: "user", Content: "last"},
|
||||
}
|
||||
kept, keptIdx, _ := Fit(msgs, 100)
|
||||
if !slices.Equal(keptIdx, []int{2}) {
|
||||
t.Fatalf("keptIdx = %v, want [2] (only the last user kept)", keptIdx)
|
||||
}
|
||||
if len(kept) != 1 || kept[0].Content != "last" {
|
||||
t.Fatalf("kept = %+v, want the last user message", kept)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_SystemOnlyMessages(t *testing.T) {
|
||||
sys1 := strings.Repeat("s ", 800)
|
||||
sys2 := strings.Repeat("t ", 200)
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: sys1},
|
||||
{Role: "system", Content: sys2},
|
||||
}
|
||||
const budget = 500
|
||||
|
||||
kept, keptIdx, count := Fit(msgs, budget)
|
||||
if count == 0 {
|
||||
t.Fatalf("Fit returned count 0, want > 0")
|
||||
}
|
||||
if len(kept) != 2 || !slices.Equal(keptIdx, []int{0, 1}) {
|
||||
t.Fatalf("got kept=%+v keptIdx=%v, want both systems", kept, keptIdx)
|
||||
}
|
||||
total := tokenizer.NumTokensFromString(kept[0].Content) + tokenizer.NumTokensFromString(kept[1].Content)
|
||||
if total > budget {
|
||||
t.Errorf("fitted total %d exceeds budget %d", total, budget)
|
||||
}
|
||||
fitted0 := tokenizer.NumTokensFromString(kept[0].Content)
|
||||
fitted1 := tokenizer.NumTokensFromString(kept[1].Content)
|
||||
if fitted0 == 0 || fitted1 == 0 {
|
||||
t.Errorf("a system message was emptied by fitting: %+v", kept)
|
||||
}
|
||||
if fitted0 >= tokenizer.NumTokensFromString(sys1) || fitted1 >= tokenizer.NumTokensFromString(sys2) {
|
||||
t.Errorf("system messages not trimmed: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_TrimsAllSystemMessages(t *testing.T) {
|
||||
sys1 := strings.Repeat("s ", 800)
|
||||
sys2 := strings.Repeat("t ", 200)
|
||||
last := "last"
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: sys1},
|
||||
{Role: "system", Content: sys2},
|
||||
{Role: "user", Content: last},
|
||||
}
|
||||
const budget = 500
|
||||
|
||||
kept, keptIdx, count := Fit(msgs, budget)
|
||||
if count == 0 {
|
||||
t.Fatalf("Fit returned count 0, want > 0")
|
||||
}
|
||||
if !slices.Equal(keptIdx, []int{0, 1, 2}) {
|
||||
t.Fatalf("keptIdx = %v, want [0 1 2]", keptIdx)
|
||||
}
|
||||
if kept[2].Content != last {
|
||||
t.Errorf("last user message not preserved: %q", kept[2].Content)
|
||||
}
|
||||
total := tokenizer.NumTokensFromString(kept[0].Content) +
|
||||
tokenizer.NumTokensFromString(kept[1].Content) +
|
||||
tokenizer.NumTokensFromString(kept[2].Content)
|
||||
if total > budget {
|
||||
t.Errorf("fitted total %d exceeds budget %d", total, budget)
|
||||
}
|
||||
fitted0 := tokenizer.NumTokensFromString(kept[0].Content)
|
||||
fitted1 := tokenizer.NumTokensFromString(kept[1].Content)
|
||||
if fitted0 == 0 || fitted1 == 0 {
|
||||
t.Errorf("a system message was emptied by fitting: %+v", kept)
|
||||
}
|
||||
if fitted0 >= tokenizer.NumTokensFromString(sys1) || fitted1 >= tokenizer.NumTokensFromString(sys2) {
|
||||
t.Errorf("system messages not trimmed: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_Step3_SystemDominates(t *testing.T) {
|
||||
// System takes >80% of tokens → preserve user, trim system.
|
||||
sysContent := strings.Repeat("a ", 800) // dominates
|
||||
userContent := strings.Repeat("b ", 100) // small, fits entirely
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: sysContent},
|
||||
{Role: "user", Content: userContent},
|
||||
}
|
||||
const budget = 500
|
||||
|
||||
kept, keptIdx, count := Fit(msgs, budget)
|
||||
if count == 0 {
|
||||
t.Fatalf("Fit returned count 0, want > 0")
|
||||
}
|
||||
if count > budget {
|
||||
t.Errorf("fitted total %d exceeds budget %d", count, budget)
|
||||
}
|
||||
if !slices.Equal(keptIdx, []int{0, 1}) {
|
||||
t.Fatalf("keptIdx = %v, want [0 1]", keptIdx)
|
||||
}
|
||||
// User preserved verbatim; system trimmed.
|
||||
if tokenizer.NumTokensFromString(kept[1].Content) != tokenizer.NumTokensFromString(userContent) {
|
||||
t.Errorf("user message not preserved: %+v", kept)
|
||||
}
|
||||
if tokenizer.NumTokensFromString(kept[0].Content) >= tokenizer.NumTokensFromString(sysContent) {
|
||||
t.Errorf("system not trimmed: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_Step3_UserDominates(t *testing.T) {
|
||||
// User takes >20% → preserve system, trim user.
|
||||
sysContent := strings.Repeat("a ", 150) // small, fits entirely
|
||||
userContent := strings.Repeat("b ", 800) // dominates
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: sysContent},
|
||||
{Role: "user", Content: userContent},
|
||||
}
|
||||
const budget = 500
|
||||
|
||||
kept, keptIdx, count := Fit(msgs, budget)
|
||||
if count == 0 {
|
||||
t.Fatalf("Fit returned count 0, want > 0")
|
||||
}
|
||||
if count > budget {
|
||||
t.Errorf("fitted total %d exceeds budget %d", count, budget)
|
||||
}
|
||||
if !slices.Equal(keptIdx, []int{0, 1}) {
|
||||
t.Fatalf("keptIdx = %v, want [0 1]", keptIdx)
|
||||
}
|
||||
// System preserved verbatim; user trimmed.
|
||||
if tokenizer.NumTokensFromString(kept[0].Content) != tokenizer.NumTokensFromString(sysContent) {
|
||||
t.Errorf("system message not preserved: %+v", kept)
|
||||
}
|
||||
if tokenizer.NumTokensFromString(kept[1].Content) >= tokenizer.NumTokensFromString(userContent) {
|
||||
t.Errorf("user not trimmed: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFit_Step3_BudgetFilledByLast locks the boundary where the last message
|
||||
// alone fills the budget (preserved == budget): the system share collapses to
|
||||
// 0, so every retained system message is kept but trimmed to empty — it must
|
||||
// NOT be reported as dropped. Python's message_fit_in retains the system entry
|
||||
// with content "" in this case too.
|
||||
func TestFit_Step3_BudgetFilledByLast(t *testing.T) {
|
||||
sysContent := strings.Repeat("a ", 3000) // dominates (>80% of tokens)
|
||||
userContent := strings.Repeat("b ", 600) // alone exceeds the budget
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: sysContent},
|
||||
{Role: "user", Content: userContent},
|
||||
}
|
||||
const budget = 500
|
||||
|
||||
kept, keptIdx, count := Fit(msgs, budget)
|
||||
if !slices.Equal(keptIdx, []int{0, 1}) {
|
||||
t.Fatalf("keptIdx = %v, want [0 1] (both messages still kept)", keptIdx)
|
||||
}
|
||||
if count > budget {
|
||||
t.Errorf("fitted total %d exceeds budget %d", count, budget)
|
||||
}
|
||||
if kept[0].Content != "" {
|
||||
t.Errorf("system message not trimmed to empty when the last message fills the budget: %+v", kept)
|
||||
}
|
||||
if tokenizer.NumTokensFromString(kept[1].Content) > budget {
|
||||
t.Errorf("user message exceeds budget after trim: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_SingleMessage(t *testing.T) {
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: strings.Repeat("x ", 1000)},
|
||||
}
|
||||
kept, keptIdx, count := Fit(msgs, 100)
|
||||
if count == 0 {
|
||||
t.Fatalf("Fit returned count 0, want > 0")
|
||||
}
|
||||
if !slices.Equal(keptIdx, []int{0}) {
|
||||
t.Fatalf("keptIdx = %v, want [0]", keptIdx)
|
||||
}
|
||||
if tokenizer.NumTokensFromString(kept[0].Content) > 100 {
|
||||
t.Errorf("single message not trimmed: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_ZeroBudget(t *testing.T) {
|
||||
// budget <= 0 should use 8192 default and not panic.
|
||||
msgs := []Message{
|
||||
{Role: "system", Content: "hello"},
|
||||
{Role: "user", Content: "world"},
|
||||
}
|
||||
kept, keptIdx, count := Fit(msgs, 0)
|
||||
if count == 0 {
|
||||
t.Fatalf("Fit returned count 0, want > 0")
|
||||
}
|
||||
if len(kept) != 2 || !slices.Equal(keptIdx, []int{0, 1}) {
|
||||
t.Fatalf("got kept=%+v keptIdx=%v, want both messages", kept, keptIdx)
|
||||
}
|
||||
if kept[0].Content != "hello" || kept[1].Content != "world" {
|
||||
t.Errorf("messages modified with default budget: %+v", kept)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFit_Empty(t *testing.T) {
|
||||
if kept, keptIdx, count := Fit(nil, 1000); kept != nil || keptIdx != nil || count != 0 {
|
||||
t.Errorf("Fit(nil) = %v, %v, %d; want nil, nil, 0", kept, keptIdx, count)
|
||||
}
|
||||
if kept, keptIdx, count := Fit([]Message{}, 1000); len(kept) != 0 || len(keptIdx) != 0 || count != 0 {
|
||||
t.Errorf("Fit(empty) = %v, %v, %d; want 0, 0, 0", kept, keptIdx, count)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user