mirror of https://github.com/chaitin/PandaWiki.git
Compare commits
2 Commits
9c02e5d93b
...
5078e93a4a
| Author | SHA1 | Date |
|---|---|---|
|
|
5078e93a4a | |
|
|
f7c0fe273b |
|
|
@ -39,7 +39,6 @@ type LLMUsecase struct {
|
|||
const (
|
||||
summaryChunkTokenLimit = 30720 // 30KB tokens per chunk
|
||||
summaryMaxChunks = 4 // max chunks to process for summary
|
||||
summaryAggregateLimit = 8192 // max tokens for aggregating summaries
|
||||
)
|
||||
|
||||
func NewLLMUsecase(config *config.Config, rag rag.RAGService, conversationRepo *pg.ConversationRepository, kbRepo *pg.KnowledgeBaseRepository, nodeRepo *pg.NodeRepository, modelRepo *pg.ModelRepository, promptRepo *pg.PromptRepo, logger *log.Logger) *LLMUsecase {
|
||||
|
|
@ -227,39 +226,7 @@ func (u *LLMUsecase) SummaryNode(ctx context.Context, model *domain.Model, name,
|
|||
return "", fmt.Errorf("failed to generate summary for document %s", name)
|
||||
}
|
||||
|
||||
// Iteratively aggregate summaries if needed
|
||||
for len(summaries) > 1 {
|
||||
joined := strings.Join(summaries, "\n\n")
|
||||
tokens, err := u.countTokens(joined)
|
||||
if err != nil {
|
||||
u.logger.Warn("Failed to count tokens for aggregation, proceeding anyway", log.Error(err))
|
||||
break
|
||||
}
|
||||
if tokens <= summaryAggregateLimit {
|
||||
break
|
||||
}
|
||||
// If still too large, aggregate in batches
|
||||
u.logger.Debug("aggregating summaries in batches", log.Int("current_summaries", len(summaries)), log.Int("tokens", tokens))
|
||||
batchSize := 2
|
||||
newSummaries := make([]string, 0, (len(summaries)+batchSize-1)/batchSize)
|
||||
for i := 0; i < len(summaries); i += batchSize {
|
||||
end := i + batchSize
|
||||
if end > len(summaries) {
|
||||
end = len(summaries)
|
||||
}
|
||||
batch := strings.Join(summaries[i:end], "\n\n")
|
||||
summary, err := u.requestSummary(ctx, chatModel, name, batch)
|
||||
if err != nil {
|
||||
u.logger.Error("Failed to aggregate summary batch", log.Int("batch_start", i), log.Error(err))
|
||||
// Fallback: use the first summary in the batch
|
||||
newSummaries = append(newSummaries, summaries[i])
|
||||
continue
|
||||
}
|
||||
newSummaries = append(newSummaries, summary)
|
||||
}
|
||||
summaries = newSummaries
|
||||
}
|
||||
|
||||
// Join all summaries and generate final summary
|
||||
joined := strings.Join(summaries, "\n\n")
|
||||
finalSummary, err := u.requestSummary(ctx, chatModel, name, joined)
|
||||
if err != nil {
|
||||
|
|
@ -301,15 +268,6 @@ func (u *LLMUsecase) requestSummary(ctx context.Context, chatModel model.BaseCha
|
|||
return strings.TrimSpace(u.trimThinking(summary)), nil
|
||||
}
|
||||
|
||||
func (u *LLMUsecase) countTokens(text string) (int, error) {
|
||||
encoding, err := tiktoken.GetEncoding("cl100k_base")
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to get encoding: %w", err)
|
||||
}
|
||||
tokens := encoding.Encode(text, nil, nil)
|
||||
return len(tokens), nil
|
||||
}
|
||||
|
||||
func (u *LLMUsecase) SplitByTokenLimit(text string, maxTokens int) ([]string, error) {
|
||||
if maxTokens <= 0 {
|
||||
return nil, fmt.Errorf("maxTokens must be greater than 0")
|
||||
|
|
|
|||
Loading…
Reference in New Issue