跳转至

Token 预算控制 (Token Budget)

Token 预算控制系统，使用 tiktoken 进行精确的 BPE Token 计数。在 QA 管道中按固定顺序执行：实体截断 → 关系截断 → 序列化统计 → 构建预热系统提示词并计数 → 计算文本块剩余预算 → 按前缀截断文本块。异步包装函数（atruncate_* / acompute_*）仅用于管道兼容性，核心计算为纯同步函数。

lightrag_langchain.token_budget.truncate_entities_by_tokens 

truncate_entities_by_tokens(
    entities: list[dict[str, Any]],
    max_tokens: int,
    model: str = "gpt-4o-mini",
) -> list[dict[str, Any]]

返回实体列表的累积 token 计数在限制内的前缀。

按顺序遍历实体，通过 tiktoken 序列化每个实体并累积 token 计数。 在第一个会导致累积计数超过 max_tokens 的实体处停止，并返回到该条目 （但不包括）为止的切片。

参数：

名称 类型 描述 默认
entities list[dict[str, Any]]

实体字典（entity_name、content、source_id 等）。

 必需
max_tokens int

序列化列表的硬性 token 上限。

 必需
model str

tiktoken 模型名称（默认 "gpt-4o-mini"）。

 'gpt-4o-mini'

返回：

类型 描述
list[dict[str, Any]]

entities 中序列化后 token 计数 <= max_tokens 的前缀。
list[dict[str, Any]]

当 max_tokens <= 0 或第一个实体单独就已超过限制时，返回空列表
list[dict[str, Any]]

（不返回部分实体）。
Example 
from lightrag_langchain.token_budget import truncate_entities_by_tokens

entities = [{"entity_name": "东莞", "content": "..."}]
truncated = truncate_entities_by_tokens(entities, max_tokens=6000)
print(len(truncated))
源代码位于： src/lightrag_langchain/token_budget.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def truncate_entities_by_tokens(
    entities: list[dict[str, Any]],
    max_tokens: int,
    model: str = "gpt-4o-mini",
) -> list[dict[str, Any]]:
    """返回实体列表的累积 token 计数在限制内的前缀。

    按顺序遍历实体，通过 tiktoken 序列化每个实体并累积 token 计数。
    在第一个会导致累积计数超过 *max_tokens* 的实体处停止，并返回到该条目
    （但不包括）为止的切片。

    Args:
        entities: 实体字典（entity_name、content、source_id 等）。
        max_tokens: 序列化列表的硬性 token 上限。
        model: tiktoken 模型名称（默认 ``"gpt-4o-mini"``）。

    Returns:
        *entities* 中序列化后 token 计数 <= *max_tokens* 的前缀。
        当 *max_tokens* <= 0 或第一个实体单独就已超过限制时，返回空列表
        （不返回部分实体）。

    Example:
        ```python
        from lightrag_langchain.token_budget import truncate_entities_by_tokens

        entities = [{"entity_name": "东莞", "content": "..."}]
        truncated = truncate_entities_by_tokens(entities, max_tokens=6000)
        print(len(truncated))
        ```
    """
    if max_tokens <= 0:
        return []

    enc = _get_tokenizer(model)
    cumulative = 0

    for i, entity in enumerate(entities):
        serialized = _serialize_item(entity)
        cumulative += len(enc.encode(serialized))
        if cumulative > max_tokens:
            return entities[:i]

    return entities

lightrag_langchain.token_budget.truncate_relations_by_tokens 

truncate_relations_by_tokens(
    relations: list[dict[str, Any]],
    max_tokens: int,
    model: str = "gpt-4o-mini",
) -> list[dict[str, Any]]

返回关系列表的累积 token 计数在限制内的前缀。

算法与 :func:truncate_entities_by_tokens 相同，但操作对象为关系字典。 作为独立函数存在，以便调用者在 Phase 4/6 pipeline 中清晰区分实体与关系上下文。

参数：

名称 类型 描述 默认
relations list[dict[str, Any]]

关系字典（src_id、tgt_id、content、description 等）。

 必需
max_tokens int

序列化列表的硬性 token 上限。

 必需
model str

tiktoken 模型名称（默认 "gpt-4o-mini"）。

 'gpt-4o-mini'

返回：

类型 描述
list[dict[str, Any]]

relations 中序列化后 token 计数 <= max_tokens 的前缀。
list[dict[str, Any]]

当 max_tokens <= 0 或第一个关系单独就已超过限制时，返回空列表。
Example 
from lightrag_langchain.token_budget import truncate_relations_by_tokens

relations = [{"src_id": "1", "tgt_id": "2", "description": "..."}]
truncated = truncate_relations_by_tokens(relations, max_tokens=8000)
print(len(truncated))
源代码位于： src/lightrag_langchain/token_budget.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def truncate_relations_by_tokens(
    relations: list[dict[str, Any]],
    max_tokens: int,
    model: str = "gpt-4o-mini",
) -> list[dict[str, Any]]:
    """返回关系列表的累积 token 计数在限制内的前缀。

    算法与 :func:`truncate_entities_by_tokens` 相同，但操作对象为关系字典。
    作为独立函数存在，以便调用者在 Phase 4/6 pipeline 中清晰区分实体与关系上下文。

    Args:
        relations: 关系字典（src_id、tgt_id、content、description 等）。
        max_tokens: 序列化列表的硬性 token 上限。
        model: tiktoken 模型名称（默认 ``"gpt-4o-mini"``）。

    Returns:
        *relations* 中序列化后 token 计数 <= *max_tokens* 的前缀。
        当 *max_tokens* <= 0 或第一个关系单独就已超过限制时，返回空列表。

    Example:
        ```python
        from lightrag_langchain.token_budget import truncate_relations_by_tokens

        relations = [{"src_id": "1", "tgt_id": "2", "description": "..."}]
        truncated = truncate_relations_by_tokens(relations, max_tokens=8000)
        print(len(truncated))
        ```
    """
    if max_tokens <= 0:
        return []

    enc = _get_tokenizer(model)
    cumulative = 0

    for i, relation in enumerate(relations):
        serialized = _serialize_item(relation)
        cumulative += len(enc.encode(serialized))
        if cumulative > max_tokens:
            return relations[:i]

    return relations

lightrag_langchain.token_budget.compute_chunk_token_budget 

compute_chunk_token_budget(
    total_tokens: int,
    sys_prompt_tokens: int,
    query_tokens: int,
    entity_tokens_used: int,
    relation_tokens_used: int,
    buffer_tokens: int = 200,
) -> int

计算 chunk 内容的剩余 token 容量。

公式（匹配上游 LightRAG 顺序）： remaining = total_tokens - sys_prompt_tokens - query_tokens - entity_tokens_used - relation_tokens_used - buffer_tokens

参数：

名称 类型 描述 默认
total_tokens int

最大总 token 数（来自 QueryParamsConfig）。

 必需
sys_prompt_tokens int

系统 prompt 的 token 数。

 必需
query_tokens int

用户查询的 token 数。

 必需
entity_tokens_used int

截断后实体列表消耗的 token 数。

 必需
relation_tokens_used int

截断后关系列表消耗的 token 数。

 必需
buffer_tokens int

用于 prompt 格式化开销的安全缓冲 （默认 200，匹配上游 LightRAG）。

 200

返回：

类型 描述
int

非负整数——chunk 内容的剩余 token 预算。
int

当预算耗尽时返回 0（绝不返回负数）。
Example 
from lightrag_langchain.token_budget import compute_chunk_token_budget

budget = compute_chunk_token_budget(
    total_tokens=30000,
    sys_prompt_tokens=500,
    query_tokens=50,
    entity_tokens_used=4000,
    relation_tokens_used=6000,
)
print(f"Chunk budget: {budget} tokens")
源代码位于： src/lightrag_langchain/token_budget.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def compute_chunk_token_budget(
    total_tokens: int,
    sys_prompt_tokens: int,
    query_tokens: int,
    entity_tokens_used: int,
    relation_tokens_used: int,
    buffer_tokens: int = 200,
) -> int:
    """计算 chunk 内容的剩余 token 容量。

    公式（匹配上游 LightRAG 顺序）：
      remaining = total_tokens
                - sys_prompt_tokens
                - query_tokens
                - entity_tokens_used
                - relation_tokens_used
                - buffer_tokens

    Args:
        total_tokens: 最大总 token 数（来自 QueryParamsConfig）。
        sys_prompt_tokens: 系统 prompt 的 token 数。
        query_tokens: 用户查询的 token 数。
        entity_tokens_used: 截断后实体列表消耗的 token 数。
        relation_tokens_used: 截断后关系列表消耗的 token 数。
        buffer_tokens: 用于 prompt 格式化开销的安全缓冲
            （默认 200，匹配上游 LightRAG）。

    Returns:
        非负整数——chunk 内容的剩余 token 预算。
        当预算耗尽时返回 0（绝不返回负数）。

    Example:
        ```python
        from lightrag_langchain.token_budget import compute_chunk_token_budget

        budget = compute_chunk_token_budget(
            total_tokens=30000,
            sys_prompt_tokens=500,
            query_tokens=50,
            entity_tokens_used=4000,
            relation_tokens_used=6000,
        )
        print(f"Chunk budget: {budget} tokens")
        ```
    """
    kg_tokens = entity_tokens_used + relation_tokens_used
    remaining = total_tokens - (sys_prompt_tokens + query_tokens + kg_tokens + buffer_tokens)
    return max(0, remaining)