74 lines
2.8 KiB
Go
74 lines
2.8 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"path/filepath"
|
||
|
|
"strconv"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/agent"
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Every prompt that can produce a window summary must state the character budget the
|
||
|
|
// validator enforces, because a model that is never told the limit overshoots it and
|
||
|
|
// parks the document.
|
||
|
|
func TestWikiSummaryPromptsStateTheEnforcedCharacterBudget(t *testing.T) {
|
||
|
|
limit := strconv.Itoa(wikiReadSummaryRunes)
|
||
|
|
for name, prompt := range map[string]string{
|
||
|
|
"read_section": wikiReadSectionPrompt,
|
||
|
|
"read_merge": wikiReadMergePrompt,
|
||
|
|
"document_page": agent.WikiSummaryPrompt,
|
||
|
|
} {
|
||
|
|
require.Contains(t, prompt, "3000", name+": a soft target below the hard limit")
|
||
|
|
require.Contains(t, prompt, limit, name+": the hard limit the validator enforces")
|
||
|
|
require.NotContains(t, prompt, "500-1500 words", name+": word counts do not bound CJK output")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func wikiSummaryShortNote() string {
|
||
|
|
return "SUMMARY: short note\n\n## Topic\nOne line per topic.\n\n## Takeaways\n- Alpha\n"
|
||
|
|
}
|
||
|
|
|
||
|
|
func wikiSummaryOverlongNote() string {
|
||
|
|
return "SUMMARY: overlong note\n\n" + strings.Repeat("Alpha Beta Gamma detail. ", 400)
|
||
|
|
}
|
||
|
|
|
||
|
|
// The window summary path must recover from a single overshoot by compressing, so the
|
||
|
|
// document is not parked for a budget the model was already told.
|
||
|
|
func TestWikiFullReadSummaryCompressionRetryCompletesTheDocument(t *testing.T) {
|
||
|
|
service := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "summary-compress"}
|
||
|
|
batch := &WikiBatchContext{}
|
||
|
|
|
||
|
|
compressions, calls := 0, 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
input := messages[len(messages)-1].Content
|
||
|
|
if strings.Contains(input, "<compression_request>") {
|
||
|
|
compressions++
|
||
|
|
return &types.ChatResponse{Content: wikiSummaryShortNote()}, nil
|
||
|
|
}
|
||
|
|
if calls == 1 {
|
||
|
|
return &types.ChatResponse{Content: wikiSummaryOverlongNote()}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: wikiSummaryShortNote()}, nil
|
||
|
|
}}
|
||
|
|
|
||
|
|
var summary string
|
||
|
|
var err error
|
||
|
|
for attempt := 0; attempt < 12; attempt++ {
|
||
|
|
summary, err = service.generateWikiFullSummary(context.Background(), model, wikiFullReadTestSource(2), "English", "", batch, scope)
|
||
|
|
if err == nil || !strings.Contains(err.Error(), "continuation required") {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
require.NoError(t, err, "one overshoot must not park the document")
|
||
|
|
require.Contains(t, summary, "short note")
|
||
|
|
require.Equal(t, 1, compressions, "exactly one compression request for the over-long segment")
|
||
|
|
require.Less(t, calls, 12, "the retry must be bounded by the per-turn operation budget")
|
||
|
|
}
|