随着应用程序日趋复杂,提示很少只是一个简单的文本字符串。通常,您会通过结合多条信息来构建提示:用于指导模型行为的系统消息、用户的当前查询、对话历史,以及可能从知识库中获取的一些文档,用于检索增强生成(RAG)系统。管理这些组件的总大小是一个不小的工程难题。令牌预算这一想法在这里显得尤为实用。令牌预算是您可以分配给提示的总令牌数量的预设上限。通过将上下文窗口视作预算,您可以编程决定如何在提示的不同部分“花费”可用令牌,从而确保在充分利用可用空间的同时,不会超出模型的限制。可视化令牌预算假设您有一个4096令牌的上下文窗口。在您开始添加内容之前,一个好的做法是为模型的响应预留一部分预算。如果您将整个窗口用于输入,模型将没有空间生成答案。一种常见做法是为输出预留总预算的25-50%。剩余的预算随后分配给您的输入组件。digraph G { rankdir=TB; splines=ortho; node [shape=box, style="rounded,filled", fillcolor="#e9ecef", fontname="Arial", fontsize=10]; edge [fontname="Arial", fontsize=9]; subgraph cluster_total { label="总上下文窗口(例如,4096令牌)"; bgcolor="#f8f9fa"; subgraph cluster_input { label="输入预算(例如,3072令牌)"; bgcolor="#dee2e6"; system [label="系统提示\n(固定成本)", fillcolor="#a5d8ff"]; query [label="用户查询\n(可变成本)", fillcolor="#bac8ff"]; context [label="检索到的文档\n(动态填充)", fillcolor="#b2f2bb"]; } output [label="为输出预留\n(例如,1024令牌)", fillcolor="#ffc9c9"]; } system -> query [style=invis]; query -> context [style=invis]; }提示的不同部分会消耗固定预算中的令牌。总上下文窗口的一部分会为模型生成的输出预留。实现预算管理器为了系统地管理令牌分配,我们可以使用一个帮助类来跟踪使用情况。TokenBudgetManager 帮助您监控已“花费”的令牌数量以及剩余数量。让我们首先为一个具有4096令牌上下文窗口的模型初始化一个管理器。from kerb.tokenizer import count_tokens, Tokenizer, truncate_to_token_limit from typing import List, Dict, Optional from datetime import datetime # 这是一个为示例提供的类。 # 在实际库中,它会直接导入。 class TokenBudgetManager: def __init__(self, budget_limit: int): self.budget_limit = budget_limit self.tokens_used = 0 def check_budget(self, tokens_needed: int) -> bool: return self.tokens_used + tokens_needed <= self.budget_limit def use_tokens(self, tokens: int) -> bool: if not self.check_budget(tokens): return False self.tokens_used += tokens return True def get_remaining_budget(self) -> int: return self.budget_limit - self.tokens_used # --- 示例用法 --- # 像gpt-3.5-turbo这样的模型的总上下文窗口 TOTAL_CONTEXT_WINDOW = 4096 # 为模型的输出预留1024令牌 RESERVED_FOR_OUTPUT = 1024 INPUT_BUDGET = TOTAL_CONTEXT_WINDOW - RESERVED_FOR_OUTPUT # 初始化输入提示的预算管理器 budget = TokenBudgetManager(budget_limit=INPUT_BUDGET) print(f"总输入预算: {budget.budget_limit} 令牌")在预算内构建RAG提示现在,我们来为一个RAG应用程序构建一个复杂的提示。我们的提示将包含一个系统消息、一个用户查询和一份检索到的文档列表。首先,我们添加固定组件并更新预算。system_prompt = "You are an AI assistant. Use the provided documents to answer the user's question." user_query = "What are the main strategies for effective token management in LLM applications?" # 计算并使用系统提示的令牌 system_tokens = count_tokens(system_prompt, tokenizer=Tokenizer.CL100K_BASE) budget.use_tokens(system_tokens) print(f"系统提示使用了 {system_tokens} 令牌。剩余: {budget.get_remaining_budget()}") # 计算并使用用户查询的令牌 query_tokens = count_tokens(user_query, tokenizer=Tokenizer.CL100K_BASE) budget.use_tokens(query_tokens) print(f"用户查询使用了 {query_tokens} 令牌。剩余: {budget.get_remaining_budget()}") # 初始化最终上下文列表 final_prompt_context = [system_prompt, user_query]接下来,我们有一份按关联性排序的检索文档列表。我们希望尽可能多地将这些文档添加到提示中,而不超出预算。我们可以遍历这些文档,检查下一个文档是否合适,如果合适就添加它。retrieved_documents = [ "Document 1: Token counting is the first step. Use a tokenizer that matches your model, like cl100k_base for GPT models, to get an accurate count of tokens for any given text. This helps in estimating API costs and managing context windows.", "Document 2: Truncation is a common strategy. If a document is too long, you can truncate it by preserving either the beginning or the end. For summaries, keep the beginning; for recent data, keep the end.", "Document 3: A token budget helps manage complex prompts. You allocate parts of the context window to different components like the system prompt, user query, and retrieved documents, while always reserving space for the model's output.", "Document 4: Context compression techniques like summarization can reduce the token count of lengthy documents while preserving the most important information, making it easier to fit more context into the prompt.", "Document 5: For conversational history, sliding window strategies are effective. You can keep a fixed number of recent messages or a total number of tokens to ensure the conversation history doesn't grow indefinitely and exceed the context limit." ] print(f"\n正在向提示添加检索到的文档...") for i, doc in enumerate(retrieved_documents): doc_tokens = count_tokens(doc, tokenizer=Tokenizer.CL100K_BASE) # 检查此文档是否符合剩余预算 if budget.check_budget(doc_tokens): budget.use_tokens(doc_tokens) final_prompt_context.append(doc) print(f"已添加文档 {i+1} ({doc_tokens} 令牌)。剩余预算: {budget.get_remaining_budget()}") else: print(f"文档 {i+1} ({doc_tokens} 令牌) 不适合。停止。") break这个简单的循环能有效地将最相关的信息填充到提示中,使其符合我们设定的限制。处理超大文档如果一个高度相关的文档太大而无法完全放入剩余预算中怎么办?在前面的示例中,我们会直接跳过它。一个更好的方法是截断文档,以便至少可以包含一部分。truncate_to_token_limit 函数非常适合此目的。让我们修改我们的逻辑以包含此步骤。# 为新示例重置预算 budget = TokenBudgetManager(budget_limit=150) # 一个较小的预算用于演示 budget.use_tokens(count_tokens(system_prompt)) budget.use_tokens(count_tokens(user_query)) final_prompt_context_truncated = [system_prompt, user_query] print(f"\n--- 截断示例 ---") print(f"起始预算: {budget.budget_limit}, 提示后剩余: {budget.get_remaining_budget()}") for i, doc in enumerate(retrieved_documents): doc_tokens = count_tokens(doc, tokenizer=Tokenizer.CL100K_BASE) remaining_budget = budget.get_remaining_budget() if remaining_budget <= 0: print("没有剩余预算。停止。") break if doc_tokens > remaining_budget: # 文档太大,所以我们截断它 print(f"文档 {i+1} ({doc_tokens} 令牌) 太大。截断以适应 {remaining_budget} 令牌。") truncated_doc = truncate_to_token_limit( doc, max_tokens=remaining_budget, tokenizer=Tokenizer.CL100K_BASE ) final_prompt_context_truncated.append(truncated_doc) # 用完剩余预算 actual_tokens = count_tokens(truncated_doc) budget.use_tokens(actual_tokens) print(f"已添加截断文档 ({actual_tokens} 令牌)。剩余预算: {budget.get_remaining_budget()}") # 没有更多空间了,所以我们停止 break else: # 文档完全适合 budget.use_tokens(doc_tokens) final_prompt_context_truncated.append(doc) print(f"已添加文档 {i+1} ({doc_tokens} 令牌)。剩余预算: {budget.get_remaining_budget()}") # 最后,将所有部分合并成一个字符串发送给LLM final_prompt = "\n\n".join(final_prompt_context_truncated)通过结合令牌计数、预算管理和动态截断,您可以构建能够自动且智能地管理LLM上下文窗口的系统。这可以防止因提示过大而导致的运行时错误,有助于控制API成本,并确保最有价值的信息始终包含在提供给模型的上下文中。