来自工业界的知识库 RAG（六），独特的 RAG 框架 dsRAG 核心亮点解读

system_prompt = """
Read the document below and extract a StructuredDocument object from it where each section of the document is centered around a single concept/topic. Whenever possible, your sections (and section titles) should match up with the natural sections of the document (i.e. Introduction, Conclusion, References, etc.). Sections can vary in length, but should generally be anywhere from a few paragraphs to a few pages long.
Each line of the document is marked with its line number in square brackets (e.g. [1], [2], [3], etc). Use the line numbers to indicate section start and end.
The start and end line numbers will be treated as inclusive. For example, if the first line of a section is line 5 and the last line is line 10, the start_index should be 5 and the end_index should be 10.
The first section must start at the first line number of the document ({start_line} in this case), and the last section must end at the last line of the document ({end_line} in this case). The sections MUST be non-overlapping and cover the entire document. In other words, they must form a partition of the document.
Section titles should be descriptive enough such that a person who is just skimming over the section titles and not actually reading the document can get a clear idea of what each section is about.
Note: the document provided to you may just be an excerpt from a larger document, rather than a complete document. Therefore, you can't always assume, for example, that the first line of the document is the beginning of the Introduction section and the last line is the end of the Conclusion section (if those section are even present).
"""

class Section(BaseModel):
    title: str = Field(description="main topic of this section of the document (very descriptive)")
    start_index: int = Field(description="line number where the section begins (inclusive)")
    end_index: int = Field(description="line number where the section ends (inclusive)")


class StructuredDocument(BaseModel):
    """obtains meaningful sections, each centered around a single concept/topic"""
    sections: List[Section] = Field(description="a list of sections of the document")

# 大模型输出结构化内容

client.chat.completions.create(
    model=model,
    response_model=StructuredDocument,
    max_tokens=4000,
    temperature=0.0,
    messages=[
        {
            "role": "system",
            "content": system_prompt.format(start_line=start_line, end_line=end_line),
        },
        {
            "role": "user",
            "content": document_with_line_numbers,
        },
    ],
)

DOCUMENT_TITLE_PROMPT = """
INSTRUCTIONS
What is the title of the following document?

Your response MUST be the title of the document, and nothing else. DO NOT respond with anything else.

{document_title_guidance}

{truncation_message}

DOCUMENT
{document_text}
""".strip()

# 计算原始评分

def get_chunk_value(chunk_info: dict, irrelevant_chunk_penalty: float, decay_rate: int):
    rank = chunk_info.get('rank', 1000)
    absolute_relevance_value = chunk_info.get('absolute_relevance_value', 0.0)
    v = np.exp(-rank / decay_rate)*absolute_relevance_value - irrelevant_chunk_penalty
    return v

# 基于文本长度进行评分缩放

def adjust_relevance_values_for_chunk_length(relevance_values: list[float], chunk_lengths: list[int], reference_length: int = 700):
    adjusted_relevance_values = []
    for relevance_value, chunk_length in zip(relevance_values, chunk_lengths):
        adjusted_relevance_values.append(relevance_value * (chunk_length / reference_length))
    return adjusted_relevance_values

for start in range(len(relevance_values)):
    # 跳过开始为 0 情况

    if relevance_values[start] < 0:
        continue

    for end in range(start+1, min(start+max_length+1, len(relevance_values)+1)):
        # 计算连续块的得分，更新可能的最大得分

        segment_value = sum(relevance_values[start:end])
        if segment_value > best_value:
            best_value = segment_value
            best_segment = (start, end)