refactor(agent): use shared message-fit in LLM component - #18092
Conversation
dosubot
Bot
added
the
size:L
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds shared token-budget fitting through ChangesMessage fitting
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant fitMessages
participant messagefit.Fit
participant schemaMessages
fitMessages->>messagefit.Fit: convert messages and apply context budget
messagefit.Fit-->>fitMessages: fitted content and token count
fitMessages->>schemaMessages: restore fitted text and preserve image-only entries
schemaMessages-->>fitMessages: validate reconstructed message sequence
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
internal/component/messagefit/messagefit.go (1)
9-14: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove migration-history references from exported documentation.
The package and
Fitcomments describe the implementation as mirroring Python code. Describe the current fitting behavior without migration references.As per coding guidelines: “keep package and doc comments aligned with the current runtime path rather than migration history.”
Also applies to: 31-43
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/component/messagefit/messagefit.go` around lines 9 - 14, Update the package documentation and the exported Fit comment to describe the current message-fitting behavior without mentioning Python, migration history, or implementation mirroring. Keep the existing descriptions of the token budget and shared caller usage accurate.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/agent/component/llm.go`:
- Around line 1065-1080: Update fitMessages so every non-empty text part in each
UserInputMultiContent is included in token fitting and correctly reconstructed,
rather than stopping at the first part; preserve non-text parts and their
ordering, and add tests covering multiple text parts in one message.
In `@internal/component/messagefit/messagefit_test.go`:
- Around line 33-60: Update TestFit_Step2_DropsMiddle to derive a budget that
accommodates the first and final long messages but excludes the middle message,
then pass that budget to Fit instead of the hardcoded 50. Remove the unused
budget assignment, and assert that msgs[1].Content is empty while the system and
final user contents remain unchanged; replace the broad non-empty checks with
these direct Step 2 assertions.
In `@internal/component/messagefit/messagefit.go`:
- Around line 55-56: Update both fit checks in Fit to use <= maxTokens so
message lists whose token count exactly equals the budget are preserved. Add a
regression test covering a full list with an intermediate message and maxTokens
equal to the full token count, asserting the intermediate message remains.
- Around line 93-113: Update Step 3 around the retained-message trimming logic
to process every system-message index in kept, not only kept[0] and
kept[len(kept)-1]. Allocate the available system-token budget across all
retained system messages, trim each via TrimContentToTokenLimit, and preserve
the existing user-message budget handling so countTokensMulti remains within
maxTokens. Add coverage for at least two retained system messages plus one user
message.
---
Nitpick comments:
In `@internal/component/messagefit/messagefit.go`:
- Around line 9-14: Update the package documentation and the exported Fit
comment to describe the current message-fitting behavior without mentioning
Python, migration history, or implementation mirroring. Keep the existing
descriptions of the token budget and shared caller usage accurate.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3f6b5e7f-f0cd-4723-80ee-78cfdbd05c57
📒 Files selected for processing (4)
internal/agent/component/llm.gointernal/agent/component/llm_test.gointernal/component/messagefit/messagefit.gointernal/component/messagefit/messagefit_test.go
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/component/messagefit/messagefit_test.go`:
- Around line 88-97: Update the assertions in the Fit test to verify that both
retained system messages have a positive token count after fitting. Compare
tokenizer.NumTokensFromString results against the original system-message token
counts rather than comparing string lengths, while preserving the existing
budget validation.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: db474d9d-d793-4791-84b0-76fd74555bd6
📒 Files selected for processing (2)
internal/component/messagefit/messagefit.gointernal/component/messagefit/messagefit_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- internal/component/messagefit/messagefit.go
|
Addressed the CodeRabbit comments (pushed to this branch):
All |
b812927 to
80d18ff
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/component/messagefit/messagefit_test.go`:
- Around line 131-147: Strengthen the Step 3 tests around Fit by snapshotting
the message content that each branch must preserve and asserting it remains
unchanged after fitting. For every affected test, assert the fitted token total
is within the requested budget; in the single-message case, calculate the result
with tokenizer.NumTokensFromString rather than byte length, while retaining the
existing trimming assertions.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 36758c99-6b79-4c29-8504-a01b565ff642
📒 Files selected for processing (4)
internal/agent/component/llm.gointernal/agent/component/llm_test.gointernal/component/messagefit/messagefit.gointernal/component/messagefit/messagefit_test.go
🚧 Files skipped from review as they are similar to previous changes (3)
- internal/component/messagefit/messagefit.go
- internal/agent/component/llm_test.go
- internal/agent/component/llm.go
| func TestFit_Step3_SystemDominates(t *testing.T) { | ||
| // System takes >80% of tokens → preserve user, trim system. | ||
| sysContent := strings.Repeat("a ", 800) // ~1600 tokens | ||
| userContent := strings.Repeat("b ", 100) // ~200 tokens | ||
| msgs := []Message{ | ||
| {Role: "system", Content: sysContent}, | ||
| {Role: "user", Content: userContent}, | ||
| } | ||
| got := Fit(msgs, 500) | ||
| if got == 0 { | ||
| t.Errorf("Fit returned 0, want > 0") | ||
| } | ||
| // System should have been trimmed, user preserved. | ||
| if len(msgs[0].Content) >= len(sysContent) { | ||
| t.Errorf("system not trimmed: got %d want < %d", len(msgs[0].Content), len(sysContent)) | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Assert the fitted token budget and preserved message.
The Step 3 tests only confirm that one message becomes shorter. They pass if Fit also modifies the message that the selected branch must preserve. The single-message test compares byte length and can pass when output still exceeds 100 tokens.
Assert that the preserved user or system content is unchanged. Assert that each fitted token total does not exceed its budget. Use tokenizer.NumTokensFromString for the single-message budget check.
Also applies to: 149-164, 166-177
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/component/messagefit/messagefit_test.go` around lines 131 - 147,
Strengthen the Step 3 tests around Fit by snapshotting the message content that
each branch must preserve and asserting it remains unchanged after fitting. For
every affected test, assert the fitted token total is within the requested
budget; in the single-message case, calculate the result with
tokenizer.NumTokensFromString rather than byte length, while retaining the
existing trimming assertions.
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).
Replace the local map-based messageFitInRaw/countAllTokens/ stringContent/setContent helpers with the shared messagefit package. Behavior is unchanged: the fitting budget still comes from the canvas max_tokens parameter.
fitMessages folded only the first non-empty text part of a UserInputMultiContent message into the token budget; any later text part survived reconstruction untrimmed and could exceed the fitted budget. Now all non-empty text parts are joined and fitted as one, the fitted text is written back to the first text part, and any additional text parts are dropped (non-text parts keep their position). Add TestFitMessages_FoldsMultipleTextParts.
messagefit.Fit now returns (kept, keptIdx, count) instead of clearing dropped entries in place. Iterate keptIdx and write kept[j].Content back to the right field; the empty-content sentinel check (and fitSource.hadText) is no longer needed because dropped entries are simply absent from keptIdx.
80d18ff to
cc3b794
Compare
The reconstruction only overwrote m.Content when kept[j].Content was non-empty. When the fitter keeps a system message but trims it to empty (the final user turn alone fills the budget), the guard left the ORIGINAL untrimmed content in place, blowing the token budget. Track whether the message originally carried text in Content and write the fitted text back unconditionally; image-only turns (no text anywhere) stay untouched.
2 participants
Summary
Switch the agent LLM component's
fitMessagesfrom the local map-basedmessageFitInRaw/countAllTokens/stringContent/setContenthelpers to the new sharedinternal/component/messagefitpackage (#18091).Why
Removes a duplicate copy of the
message_fit_instrategy now that the shared package exists; the ingestion Extractor uses the same package in #18095.Changes
fitMessagesconvertsschema.Message→messagefit.Message, callsmessagefit.Fit, then writes the trimmed text back to the right field (plainContentor the text part ofUserInputMultiContent).fitMessagesis preserved when kept by the fitter.fitSourcenow tracks whether the message originally carried text inContent.max_tokensparameter — behavior unchanged (changed in fix(agent): fit LLM prompts to model content_length #18093).Testing
bash build.sh --test ./internal/agent/component/...passes.TestFitMessages_SystemKeptButEmptied).main(7b2d052f8) with feat(component): add shared message-fit package #18091 and feat(dao): add shared ResolveModelContentLength helper #18106 merged; no carried copies ofmessagefitremain.fitMessagesto the non-mutatingmessagefit.FitAPI (kept,keptIdx,count): reconstruct fromkeptIdxinstead of filteringContent == "".Depends on #18091.