mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-01 21:37:33 +08:00
90f46b0b4d
Ports doc-level auto-metadata extraction to Go and adds the knowledge_compiler component with scheduler/routing. Fixes Extractor metadata injection type assertion and enable_metadata default-on.
89 lines
2.6 KiB
Go
89 lines
2.6 KiB
Go
package mindmap
|
|
|
|
import (
|
|
"strings"
|
|
|
|
"ragflow/internal/ingestion/component/knowledge_compiler/common"
|
|
)
|
|
|
|
// mindMapExtractionPrompt mirrors MIND_MAP_EXTRACTION_PROMPT verbatim
|
|
// (rag/graphrag/general/mind_map_prompt.py). It is the SYSTEM message; the
|
|
// user turn is the fixed string "Output:" (see userMessage).
|
|
const mindMapExtractionPrompt = `
|
|
- Role: You're a talent text processor to summarize a piece of text into a mind map.
|
|
|
|
- Step of task:
|
|
1. Generate a title for user's 'TEXT'。
|
|
2. Classify the 'TEXT' into sections of a mind map.
|
|
3. If the subject matter is really complex, split them into sub-sections and sub-subsections.
|
|
4. Add a shot content summary of the bottom level section.
|
|
|
|
- Output requirement:
|
|
- Generate at least 4 levels.
|
|
- Always try to maximize the number of sub-sections.
|
|
- In language of 'Text'
|
|
- MUST IN FORMAT OF MARKDOWN
|
|
|
|
-TEXT-
|
|
{input_text}
|
|
|
|
`
|
|
|
|
// userMessage is the fixed user turn appended after the rendered system
|
|
// prompt (Python: [{"role": "user", "content": "Output:"}]).
|
|
const userMessage = "Output:"
|
|
|
|
// mindmapMaxLength stands in for chat_mdl.max_length — the ChatInvoker seam
|
|
// does not expose the model window, so the batch budget uses the same
|
|
// conservative constant the other variants use. Python derives the budget as
|
|
// max(max_length*0.8, max_length-512).
|
|
const mindmapMaxLength = 4096
|
|
|
|
// batchBudget mirrors max(max_length*0.8, max_length-512).
|
|
func batchBudget() int {
|
|
a := mindmapMaxLength * 8 / 10
|
|
b := mindmapMaxLength - 512
|
|
if a > b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|
|
|
|
// renderPrompt mirrors perform_variable_replacements for the single
|
|
// {input_text} variable (prompt_variables is empty at the Python call site).
|
|
func renderPrompt(text string) string {
|
|
return strings.ReplaceAll(mindMapExtractionPrompt, "{input_text}", text)
|
|
}
|
|
|
|
// packSections mirrors MindMapExtractor.__call__'s batching: sections are
|
|
// concatenated WITHOUT a separator into batches whose token count stays
|
|
// under the budget; a single section is never split (an oversized section
|
|
// forms its own batch).
|
|
func packSections(sections []string, tok common.Tokenizer) []string {
|
|
budget := batchBudget()
|
|
var batches []string
|
|
var cur strings.Builder
|
|
cnt := 0
|
|
for _, s := range sections {
|
|
sc := numTokens(tok, s)
|
|
if cnt+sc >= budget && cur.Len() > 0 {
|
|
batches = append(batches, cur.String())
|
|
cur.Reset()
|
|
cnt = 0
|
|
}
|
|
cur.WriteString(s)
|
|
cnt += sc
|
|
}
|
|
if cur.Len() > 0 {
|
|
batches = append(batches, cur.String())
|
|
}
|
|
return batches
|
|
}
|
|
|
|
func numTokens(tok common.Tokenizer, s string) int {
|
|
if tok != nil {
|
|
return tok.NumTokens(s)
|
|
}
|
|
return common.EstimateTokens(s)
|
|
}
|