Truncate retrieved docs to fit within context limit. Args: docs: List of documents to truncate question: The question being asked system_message: Optional system message max_context_tokens: Maximum tokens allowed for context doc_fo
(
self,
docs: List[str],
question: str,
system_message: Optional[str] = None,
max_context_tokens: int = 120000,
doc_format: str = "[Memory {index}]\n{content}",
overhead_tokens: int = 100,
min_remaining_tokens: int = 100,
)
| 152 | } |
| 153 | |
| 154 | def truncate_docs_to_context( |
| 155 | self, |
| 156 | docs: List[str], |
| 157 | question: str, |
| 158 | system_message: Optional[str] = None, |
| 159 | max_context_tokens: int = 120000, |
| 160 | doc_format: str = "[Memory {index}]\n{content}", |
| 161 | overhead_tokens: int = 100, |
| 162 | min_remaining_tokens: int = 100, |
| 163 | ) -> List[str]: |
| 164 | """Truncate retrieved docs to fit within context limit. |
| 165 | |
| 166 | Args: |
| 167 | docs: List of documents to truncate |
| 168 | question: The question being asked |
| 169 | system_message: Optional system message |
| 170 | max_context_tokens: Maximum tokens allowed for context |
| 171 | doc_format: Format string for each doc (must contain {index} and {content}) |
| 172 | overhead_tokens: Reserved tokens for formatting overhead |
| 173 | min_remaining_tokens: Minimum tokens needed to include partial doc |
| 174 | |
| 175 | Returns: |
| 176 | List of truncated documents that fit within context limit |
| 177 | """ |
| 178 | if not docs: |
| 179 | return docs |
| 180 | |
| 181 | question_tokens = self.count_tokens(question) |
| 182 | system_tokens = self.count_tokens(system_message) if system_message else 0 |
| 183 | available_tokens = max_context_tokens - question_tokens - system_tokens - overhead_tokens |
| 184 | |
| 185 | if available_tokens <= 0: |
| 186 | return [] |
| 187 | |
| 188 | truncated_docs = [] |
| 189 | current_tokens = 0 |
| 190 | |
| 191 | for doc in docs: |
| 192 | doc_tokens = self.count_tokens(doc) |
| 193 | format_overhead = self.count_tokens( |
| 194 | doc_format.format(index=len(truncated_docs) + 1, content="") |
| 195 | ) |
| 196 | |
| 197 | if current_tokens + doc_tokens + format_overhead <= available_tokens: |
| 198 | truncated_docs.append(doc) |
| 199 | current_tokens += doc_tokens + format_overhead |
| 200 | else: |
| 201 | remaining_tokens = available_tokens - current_tokens - format_overhead |
| 202 | if remaining_tokens > min_remaining_tokens: |
| 203 | tokens = self._tokenizer.encode(doc) |
| 204 | truncated_text = self._tokenizer.decode(tokens[:remaining_tokens]) |
| 205 | truncated_docs.append(truncated_text + "...") |
| 206 | break |
| 207 | |
| 208 | return truncated_docs |