@samitouri / QOSami_Agents_Zero-1 / commits / 3dceaca6

caching and ctx window optimizations

frdel committed Feb 11, 2026 at 23:04 UTC 3dceaca64ec49d4149b9ae91a0e8e66a4692d132

7 files changed +82 -39

agent.py

@@ -798,6 +798,7 @@ class Agent:

798 response_callback: Callable[[str, str], Awaitable[None]] | None = None,

799 reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,

800 background: bool = False,

801 + explicit_caching: bool = True,

802 ):

803 response = ""

804

@@ -812,6 +813,7 @@ class Agent:

813 rate_limiter_callback=(

814 self.rate_limiter_callback if not background else None

815 ),

816 + explicit_caching=explicit_caching,

817 )

818

819 return response, reasoning

models.py

+12 -2

@@ -316,7 +316,7 @@ class LiteLLMChatWrapper(SimpleChatModel):

316 def _llm_type(self) -> str:

317 return "litellm-chat"

318

319 - def _convert_messages(self, messages: List[BaseMessage]) -> List[dict]:

319 + def _convert_messages(self, messages: List[BaseMessage], explicit_caching: bool = False) -> List[dict]:

320 result = []

321 # Map LangChain message types to LiteLLM roles

322 role_mapping = {

@@ -362,6 +362,15 @@ class LiteLLMChatWrapper(SimpleChatModel):

362 message_dict["tool_call_id"] = tool_call_id

363

364 result.append(message_dict)

365 +

366 + if explicit_caching and result:

367 + if result[0]["role"] == "system":

368 + result[0]["cache_control"] = {"type": "ephemeral"}

369 + for i in range(len(result) - 1, -1, -1):

370 + if result[i]["role"] == "assistant":

371 + result[i]["cache_control"] = {"type": "ephemeral"}

372 + break

373 +

374 return result

375

376 def _call(

@@ -464,6 +473,7 @@ class LiteLLMChatWrapper(SimpleChatModel):

473 rate_limiter_callback: (

474 Callable[[str, str, int, int], Awaitable[bool]] | None

475 ) = None,

476 + explicit_caching: bool = False,

477 **kwargs: Any,

478 ) -> Tuple[str, str]:

479

@@ -478,7 +488,7 @@ class LiteLLMChatWrapper(SimpleChatModel):

488 messages.append(HumanMessage(content=user_message))

489

490 # convert to litellm format

481 - msgs_conv = self._convert_messages(messages)

491 + msgs_conv = self._convert_messages(messages, explicit_caching=explicit_caching)

492

493 # Apply rate limiting if configured

494 limiter = await apply_rate_limiter(

prompts/fw.topic_summary.sys.md

+2 -1

@@ -6,7 +6,8 @@ You must return a single summary of all records

7 # Expected output

8 Your output will be a text of the summary

9 -Length of the text should be one paragraph, approximately 100 words

9 +Summary must be shorter than original messages

10 +Length of the text should be maximum one paragraph, approximately 100 words, shorter if original is shorter

11 No intro

12 No conclusion

13 No formatting

python/extensions/message_loop_prompts_after/_50_recall_memories.py

+5 -1

@@ -8,6 +8,7 @@ from python.helpers import dirty_json, errors, settings, log

9 DATA_NAME_TASK = "_recall_memories_task"

10 DATA_NAME_ITER = "_recall_memories_iter"

11 +SEARCH_TIMEOUT = 30

14 class RecallMemories(Extension):

@@ -38,7 +39,10 @@ class RecallMemories(Extension):

39 )

41 task = asyncio.create_task(

41 - self.search_memories(loop_data=loop_data, log_item=log_item, **kwargs)

42 + asyncio.wait_for(

43 + self.search_memories(loop_data=loop_data, log_item=log_item, **kwargs),

44 + timeout=SEARCH_TIMEOUT,

45 + )

46 )

47 else:

48 task = None

python/helpers/document_query.py

+2 -1

@@ -433,7 +433,8 @@ class DocumentQueryHelper:

433 messages=[

434 SystemMessage(content=qa_system_message),

435 HumanMessage(content=qa_user_message),

436 - ]

436 + ],

437 + explicit_caching=False,

438 )

439

440 self.progress_callback(f"Q&A process completed")

python/helpers/history.py

+58 -33

@@ -10,13 +10,16 @@ from enum import Enum

10 from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage

12 BULK_MERGE_COUNT = 3

13 -TOPICS_KEEP_COUNT = 3

13 +TOPICS_MERGE_COUNT = 3

14 CURRENT_TOPIC_RATIO = 0.5

15 HISTORY_TOPIC_RATIO = 0.3

16 HISTORY_BULK_RATIO = 0.2

17 -TOPIC_COMPRESS_RATIO = 0.65

18 -LARGE_MESSAGE_TO_TOPIC_RATIO = 0.5

17 +CURRENT_TOPIC_ATTENTION_COMPRESSION = 0.65 # compress current topic's attention window to 65% of size

18 +HISTORY_TOPIC_ATTENTION_COMPRESSION = 0 # compress history topic's attention window to 0% of size - only request and response remain intact

19 +LARGE_MESSAGE_TO_CURRENT_TOPIC_RATIO = 0.5

20 +LARGE_MESSAGE_TO_HISTORY_TOPIC_RATIO = 0.2

21 RAW_MESSAGE_OUTPUT_TEXT_TRIM = 100

22 +COMPRESSION_TARGET_RATIO = 0.8

25 class RawMessage(TypedDict):

@@ -155,13 +158,12 @@ class Topic(Record):

158 self.summary = await self.summarize_messages(self.messages)

159 return self.summary

160

158 - async def compress_large_messages(self) -> bool:

161 + def compress_large_messages(self, message_ratio: float = CURRENT_TOPIC_RATIO * LARGE_MESSAGE_TO_CURRENT_TOPIC_RATIO) -> bool:

162 set = settings.get_settings()

163 msg_max_size = (

164 set["chat_model_ctx_length"]

165 * set["chat_model_ctx_history"]

163 - * CURRENT_TOPIC_RATIO

164 - * LARGE_MESSAGE_TO_TOPIC_RATIO

166 + * message_ratio

167 )

168 large_msgs = []

169 for m in (m for m in self.messages if not m.summary):

@@ -195,27 +197,29 @@ class Topic(Record):

197 return False

198

199 async def compress(self) -> bool:

198 - compress = await self.compress_large_messages()

200 + compress = self.compress_large_messages()

201 if not compress:

202 compress = await self.compress_attention()

203 return compress

204

203 - async def compress_attention(self) -> bool:

205 + async def compress_attention(self, ratio: float = CURRENT_TOPIC_ATTENTION_COMPRESSION) -> bool:

206

205 - if len(self.messages) > 2:

206 - cnt_to_sum = math.ceil((len(self.messages) - 2) * TOPIC_COMPRESS_RATIO)

207 - msg_to_sum = self.messages[1 : cnt_to_sum + 1]

208 - summary = await self.summarize_messages(msg_to_sum)

209 - sum_msg_content = self.history.agent.parse_prompt(

210 - "fw.msg_summary.md", summary=summary

211 - )

212 - sum_msg = Message(False, sum_msg_content)

213 - self.messages[1 : cnt_to_sum + 1] = [sum_msg]

214 - return True

215 - return False

207 + middle = len(self.messages) - 2

208 + if middle < 2:

209 + return False

210 + cnt_to_sum = middle - math.floor(middle * ratio)

211 + if cnt_to_sum < 1:

212 + return False

213 + msg_to_sum = self.messages[1 : cnt_to_sum + 1]

214 + summary = await self.summarize_messages(msg_to_sum)

215 + sum_msg_content = self.history.agent.parse_prompt(

216 + "fw.msg_summary.md", summary=summary

217 + )

218 + sum_msg = Message(False, sum_msg_content)

219 + self.messages[1 : cnt_to_sum + 1] = [sum_msg]

220 + return True

221

222 async def summarize_messages(self, messages: list[Message]):

218 - # FIXME: vision bytes are sent to utility LLM, send summary instead

223 msg_txt = [m.output_text() for m in messages]

224 summary = await self.history.agent.call_utility_model(

225 system=self.history.agent.read_prompt("fw.topic_summary.sys.md"),

@@ -363,22 +367,38 @@ class History(Record):

367

368 async def compress(self):

369 compressed = False

370 + total = _get_ctx_size_for_history()

371 + curr, hist, bulk = (

372 + self.get_current_topic_tokens(),

373 + self.get_topics_tokens(),

374 + self.get_bulks_tokens(),

375 + )

376 + if (curr + hist + bulk) <= total:

377 + return False

378 +

379 + target = total * COMPRESSION_TARGET_RATIO

380 + prev_total = curr + hist + bulk

381 while True:

382 curr, hist, bulk = (

383 self.get_current_topic_tokens(),

384 self.get_topics_tokens(),

385 self.get_bulks_tokens(),

386 )

372 - total = _get_ctx_size_for_history()

387 +

388 + # safeguard against infinite loop in case LLM bloats the summary for some reason

389 + if (curr + hist + bulk) >= prev_total:

390 + break

391 + prev_total = curr + hist + bulk

392 +

393 ratios = [

394 (curr, CURRENT_TOPIC_RATIO, "current_topic"),

395 (hist, HISTORY_TOPIC_RATIO, "history_topic"),

396 (bulk, HISTORY_BULK_RATIO, "history_bulk"),

397 ]

378 - ratios = sorted(ratios, key=lambda x: (x[0] / total) / x[1], reverse=True)

398 + ratios = sorted(ratios, key=lambda x: (x[0] / target) / x[1], reverse=True)

399 compressed_part = False

400 for ratio in ratios:

381 - if ratio[0] > ratio[1] * total:

401 + if ratio[0] > ratio[1] * target:

402 over_part = ratio[2]

403 if over_part == "current_topic":

404 compressed_part = await self.current.compress()

@@ -394,24 +414,29 @@ class History(Record):

414 continue

415 else:

416 return compressed

417 + return compressed

418

419 async def compress_topics(self) -> bool:

399 - # summarize topics one by one

420 +

421 + # 1. first identify large messages and compress them cheaply

422 for topic in self.topics:

401 - if not topic.summary:

402 - await topic.summarize()

423 + if topic.compress_large_messages(HISTORY_TOPIC_RATIO*LARGE_MESSAGE_TO_HISTORY_TOPIC_RATIO):

424 return True

425

405 - # move oldest topic to bulks and summarize

426 + # 2. summarize topics attention window one by one

427 for topic in self.topics:

428 + if await topic.compress_attention(HISTORY_TOPIC_ATTENTION_COMPRESSION):

429 + return True

430 +

431 + # 3. move oldest topics to bulks in chunks

432 + if self.topics:

433 + count = TOPICS_MERGE_COUNT if len(self.topics) >= TOPICS_MERGE_COUNT else 1

434 + chunk = self.topics[:count]

435 bulk = Bulk(history=self)

408 - bulk.records.append(topic)

409 - if topic.summary:

410 - bulk.summary = topic.summary

411 - else:

412 - await bulk.summarize()

436 + bulk.records.extend(chunk)

437 + await bulk.summarize()

438 self.bulks.append(bulk)

414 - self.topics.remove(topic)

439 + self.topics[:count] = []

440 return True

441 return False

442

python/helpers/memory_consolidation.py

+1 -1

@@ -82,7 +82,7 @@ class MemoryConsolidator:

83 Args:

84 new_memory: The new memory content to process

85 - area: Memory area (MAIN, FRAGMENTS, SOLUTIONS, INSTRUMENTS)

85 + area: Memory area (MAIN, FRAGMENTS, SOLUTIONS)

86 metadata: Initial metadata for the memory

87 log_item: Optional log item for progress tracking