| 1 | from abc import abstractmethod |
| 2 | import asyncio |
| 3 | from collections import OrderedDict |
| 4 | from collections.abc import Mapping |
| 5 | import json |
| 6 | import math |
| 7 | import uuid |
| 8 | from typing import Coroutine, Literal, TypedDict, cast, Union, Dict, List, Any |
| 9 | from helpers import messages, tokens, settings, call_llm |
| 10 | from enum import Enum |
| 11 | from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage |
| 12 | from plugins._model_config.helpers.model_config import get_chat_model_config |
| 13 | |
| 14 | |
| 15 | BULK_MERGE_COUNT = 3 |
| 16 | TOPICS_MERGE_COUNT = 3 |
| 17 | CURRENT_TOPIC_RATIO = 0.5 |
| 18 | HISTORY_TOPIC_RATIO = 0.3 |
| 19 | HISTORY_BULK_RATIO = 0.2 |
| 20 | CURRENT_TOPIC_ATTENTION_COMPRESSION = 0.65 # compress current topic's attention window to 65% of size |
| 21 | HISTORY_TOPIC_ATTENTION_COMPRESSION = 0 # compress history topic's attention window to 0% of size - only request and response remain intact |
| 22 | LARGE_MESSAGE_TO_CURRENT_TOPIC_RATIO = 0.5 |
| 23 | LARGE_MESSAGE_TO_HISTORY_TOPIC_RATIO = 0.2 |
| 24 | RAW_MESSAGE_OUTPUT_TEXT_TRIM = 100 |
| 25 | COMPRESSION_TARGET_RATIO = 0.8 |
| 26 | |
| 27 | |
| 28 | class RawMessage(TypedDict): |
| 29 | raw_content: "MessageContent" |
| 30 | preview: str | None |
| 31 | |
| 32 | |
| 33 | MessageContent = Union[ |
| 34 | List["MessageContent"], |
| 35 | Dict[str, "MessageContent"], |
| 36 | List[Dict[str, "MessageContent"]], |
| 37 | str, |
| 38 | List[str], |
| 39 | RawMessage, |
| 40 | ] |
| 41 | |
| 42 | |
| 43 | class OutputMessage(TypedDict, total=False): |
| 44 | ai: bool |
| 45 | content: MessageContent |
| 46 | metadata: dict[str, Any] |
| 47 | id: str |
| 48 | sequence: int |
| 49 | |
| 50 | |
| 51 | class Record: |
| 52 | def __init__(self): |
| 53 | pass |
| 54 | |
| 55 | @abstractmethod |
| 56 | def get_tokens(self) -> int: |
| 57 | pass |
| 58 | |
| 59 | @abstractmethod |
| 60 | async def compress(self) -> bool: |
| 61 | pass |
| 62 | |
| 63 | @abstractmethod |
| 64 | def output(self) -> list[OutputMessage]: |
| 65 | pass |
| 66 | |
| 67 | @abstractmethod |
| 68 | async def summarize(self) -> str: |
| 69 | pass |
| 70 | |
| 71 | @abstractmethod |
| 72 | def to_dict(self) -> dict: |
| 73 | pass |
| 74 | |
| 75 | @staticmethod |
| 76 | def from_dict(data: dict, history: "History"): |
| 77 | cls = data["_cls"] |
| 78 | return globals()[cls].from_dict(data, history=history) |
| 79 | |
| 80 | def output_langchain(self): |
| 81 | return output_langchain(self.output()) |
| 82 | |
| 83 | def output_text(self, human_label="user", ai_label="ai"): |
| 84 | return output_text(self.output(), ai_label, human_label) |
| 85 | |
| 86 | |
| 87 | class Message(Record): |
| 88 | def __init__( |
| 89 | self, |
| 90 | ai: bool, |
| 91 | content: MessageContent, |
| 92 | tokens: int = 0, |
| 93 | id: str = "", |
| 94 | metadata: dict[str, Any] | None = None, |
| 95 | sequence: int = 0, |
| 96 | ): |
| 97 | self.id = id or str(uuid.uuid4()) |
| 98 | self.ai = ai |
| 99 | self.content = content |
| 100 | self.metadata = metadata or {} |
| 101 | self.sequence = sequence |
| 102 | self.summary: str = "" |
| 103 | self.tokens: int = tokens or self.calculate_tokens() |
| 104 | |
| 105 | def get_tokens(self) -> int: |
| 106 | if not self.tokens: |
| 107 | self.tokens = self.calculate_tokens() |
| 108 | return self.tokens |
| 109 | |
| 110 | def calculate_tokens(self): |
| 111 | text = self.output_text() |
| 112 | return tokens.approximate_tokens(text) |
| 113 | |
| 114 | def set_summary(self, summary: str): |
| 115 | self.summary = summary |
| 116 | self.tokens = self.calculate_tokens() |
| 117 | |
| 118 | async def compress(self): |
| 119 | return False |
| 120 | |
| 121 | def output(self): |
| 122 | return [ |
| 123 | OutputMessage( |
| 124 | ai=self.ai, |
| 125 | content=self.summary or self.content, |
| 126 | metadata=self.metadata, |
| 127 | id=self.id, |
| 128 | sequence=self.sequence, |
| 129 | ) |
| 130 | ] |
| 131 | |
| 132 | def output_langchain(self): |
| 133 | return output_langchain(self.output()) |
| 134 | |
| 135 | def output_text(self, human_label="user", ai_label="ai"): |
| 136 | return output_text(self.output(), ai_label, human_label) |
| 137 | |
| 138 | def to_dict(self): |
| 139 | return { |
| 140 | "_cls": "Message", |
| 141 | "id": self.id, |
| 142 | "ai": self.ai, |
| 143 | "content": self.content, |
| 144 | "metadata": self.metadata, |
| 145 | "sequence": self.sequence, |
| 146 | "summary": self.summary, |
| 147 | "tokens": self.tokens, |
| 148 | } |
| 149 | |
| 150 | @staticmethod |
| 151 | def from_dict(data: dict, history: "History"): |
| 152 | content = data.get("content", "Content lost") |
| 153 | metadata = data.get("metadata", {}) |
| 154 | metadata = metadata if isinstance(metadata, dict) else {} |
| 155 | if data["ai"]: |
| 156 | from helpers.llm_result import result_from_metadata |
| 157 | |
| 158 | result = result_from_metadata(metadata) |
| 159 | if result: |
| 160 | metadata = {**metadata, **result.metadata()} |
| 161 | msg = Message( |
| 162 | ai=data["ai"], |
| 163 | content=content, |
| 164 | id=data.get("id", ""), |
| 165 | metadata=metadata, |
| 166 | sequence=int(data.get("sequence", 0) or 0), |
| 167 | ) |
| 168 | msg.summary = data.get("summary", "") |
| 169 | msg.tokens = data.get("tokens", 0) |
| 170 | return msg |
| 171 | |
| 172 | |
| 173 | class Topic(Record): |
| 174 | def __init__(self, history: "History"): |
| 175 | self.history = history |
| 176 | self.summary: str = "" |
| 177 | self.messages: list[Message] = [] |
| 178 | |
| 179 | def get_tokens(self): |
| 180 | if self.summary: |
| 181 | return tokens.approximate_tokens(self.summary) |
| 182 | else: |
| 183 | return sum(msg.get_tokens() for msg in self.messages) |
| 184 | |
| 185 | def add_message( |
| 186 | self, |
| 187 | ai: bool, |
| 188 | content: MessageContent, |
| 189 | tokens: int = 0, |
| 190 | id: str = "", |
| 191 | metadata: dict[str, Any] | None = None, |
| 192 | sequence: int = 0, |
| 193 | ) -> Message: |
| 194 | msg = Message( |
| 195 | ai=ai, |
| 196 | content=content, |
| 197 | tokens=tokens, |
| 198 | id=id, |
| 199 | metadata=metadata, |
| 200 | sequence=sequence, |
| 201 | ) |
| 202 | self.messages.append(msg) |
| 203 | return msg |
| 204 | |
| 205 | def output(self) -> list[OutputMessage]: |
| 206 | if self.summary: |
| 207 | return [OutputMessage(ai=False, content=self.summary)] |
| 208 | else: |
| 209 | msgs = [m for r in self.messages for m in r.output()] |
| 210 | return msgs |
| 211 | |
| 212 | async def summarize(self): |
| 213 | self.summary = await self.summarize_messages(self.messages) |
| 214 | return self.summary |
| 215 | |
| 216 | def compress_large_messages(self, message_ratio: float = CURRENT_TOPIC_RATIO * LARGE_MESSAGE_TO_CURRENT_TOPIC_RATIO) -> bool: |
| 217 | from plugins._model_config.helpers.model_config import get_chat_model_config |
| 218 | chat_cfg = get_chat_model_config() |
| 219 | ctx_length = int(chat_cfg.get("ctx_length", 128000)) |
| 220 | ctx_history = float(chat_cfg.get("ctx_history", 0.7)) |
| 221 | msg_max_size = ( |
| 222 | ctx_length |
| 223 | * ctx_history |
| 224 | * message_ratio |
| 225 | ) |
| 226 | large_msgs = [] |
| 227 | for m in (m for m in self.messages if not m.summary): |
| 228 | # TODO refactor this |
| 229 | out = m.output() |
| 230 | text = output_text(out) |
| 231 | tok = m.get_tokens() |
| 232 | leng = len(text) |
| 233 | if tok > msg_max_size: |
| 234 | large_msgs.append((m, tok, leng, out)) |
| 235 | large_msgs.sort(key=lambda x: x[1], reverse=True) |
| 236 | for msg, tok, leng, out in large_msgs: |
| 237 | trim_to_chars = leng * (msg_max_size / tok) |
| 238 | # raw messages will be replaced as a whole, they would become invalid when truncated |
| 239 | if _is_raw_message(out[0]["content"]): |
| 240 | msg.set_summary( |
| 241 | "Message content replaced to save space in context window" |
| 242 | ) |
| 243 | |
| 244 | # regular messages will be truncated |
| 245 | else: |
| 246 | trunc = messages.truncate_dict_by_ratio( |
| 247 | self.history.agent, |
| 248 | out[0]["content"], |
| 249 | trim_to_chars * 1.15, |
| 250 | trim_to_chars * 0.85, |
| 251 | ) |
| 252 | msg.set_summary(_json_dumps(trunc)) |
| 253 | |
| 254 | return True |
| 255 | return False |
| 256 | |
| 257 | async def compress(self) -> bool: |
| 258 | compress = self.compress_large_messages() |
| 259 | if not compress: |
| 260 | compress = await self.compress_attention() |
| 261 | return compress |
| 262 | |
| 263 | async def compress_attention(self, ratio: float = CURRENT_TOPIC_ATTENTION_COMPRESSION) -> bool: |
| 264 | |
| 265 | middle = len(self.messages) - 2 |
| 266 | if middle < 2: |
| 267 | return False |
| 268 | cnt_to_sum = middle - math.floor(middle * ratio) |
| 269 | if cnt_to_sum < 1: |
| 270 | return False |
| 271 | msg_to_sum = self.messages[1 : cnt_to_sum + 1] |
| 272 | summary = await self.summarize_messages(msg_to_sum) |
| 273 | sum_msg_content = self.history.agent.parse_prompt( |
| 274 | "fw.msg_summary.md", summary=summary |
| 275 | ) |
| 276 | sum_msg = Message(False, sum_msg_content) |
| 277 | self.messages[1 : cnt_to_sum + 1] = [sum_msg] |
| 278 | return True |
| 279 | |
| 280 | async def summarize_messages(self, messages: list[Message]): |
| 281 | msg_txt = [m.output_text() for m in messages] |
| 282 | summary = await self.history.agent.call_utility_model( |
| 283 | system=self.history.agent.read_prompt("fw.topic_summary.sys.md"), |
| 284 | message=self.history.agent.read_prompt( |
| 285 | "fw.topic_summary.msg.md", content=msg_txt |
| 286 | ), |
| 287 | ) |
| 288 | return summary |
| 289 | |
| 290 | def to_dict(self): |
| 291 | return { |
| 292 | "_cls": "Topic", |
| 293 | "summary": self.summary, |
| 294 | "messages": [m.to_dict() for m in self.messages], |
| 295 | } |
| 296 | |
| 297 | @staticmethod |
| 298 | def from_dict(data: dict, history: "History"): |
| 299 | topic = Topic(history=history) |
| 300 | topic.summary = data.get("summary", "") |
| 301 | topic.messages = [ |
| 302 | Message.from_dict(m, history=history) for m in data.get("messages", []) |
| 303 | ] |
| 304 | return topic |
| 305 | |
| 306 | |
| 307 | class Bulk(Record): |
| 308 | def __init__(self, history: "History"): |
| 309 | self.history = history |
| 310 | self.summary: str = "" |
| 311 | self.records: list[Record] = [] |
| 312 | |
| 313 | def get_tokens(self): |
| 314 | if self.summary: |
| 315 | return tokens.approximate_tokens(self.summary) |
| 316 | else: |
| 317 | return sum([r.get_tokens() for r in self.records]) |
| 318 | |
| 319 | def output( |
| 320 | self, human_label: str = "user", ai_label: str = "ai" |
| 321 | ) -> list[OutputMessage]: |
| 322 | if self.summary: |
| 323 | return [OutputMessage(ai=False, content=self.summary)] |
| 324 | else: |
| 325 | msgs = [m for r in self.records for m in r.output()] |
| 326 | return msgs |
| 327 | |
| 328 | async def compress(self): |
| 329 | return False |
| 330 | |
| 331 | async def summarize(self): |
| 332 | self.summary = await self.history.agent.call_utility_model( |
| 333 | system=self.history.agent.read_prompt("fw.topic_summary.sys.md"), |
| 334 | message=self.history.agent.read_prompt( |
| 335 | "fw.topic_summary.msg.md", content=self.output_text() |
| 336 | ), |
| 337 | ) |
| 338 | return self.summary |
| 339 | |
| 340 | def to_dict(self): |
| 341 | return { |
| 342 | "_cls": "Bulk", |
| 343 | "summary": self.summary, |
| 344 | "records": [r.to_dict() for r in self.records], |
| 345 | } |
| 346 | |
| 347 | @staticmethod |
| 348 | def from_dict(data: dict, history: "History"): |
| 349 | bulk = Bulk(history=history) |
| 350 | bulk.summary = data["summary"] |
| 351 | cls = data["_cls"] |
| 352 | bulk.records = [Record.from_dict(r, history=history) for r in data["records"]] |
| 353 | return bulk |
| 354 | |
| 355 | |
| 356 | class History(Record): |
| 357 | def __init__(self, agent): |
| 358 | from agent import Agent |
| 359 | |
| 360 | self.counter = 0 |
| 361 | self.bulks: list[Bulk] = [] |
| 362 | self.topics: list[Topic] = [] |
| 363 | self.current = Topic(history=self) |
| 364 | self.agent: Agent = agent |
| 365 | |
| 366 | def get_tokens(self) -> int: |
| 367 | return ( |
| 368 | self.get_bulks_tokens() |
| 369 | + self.get_topics_tokens() |
| 370 | + self.get_current_topic_tokens() |
| 371 | ) |
| 372 | |
| 373 | def is_over_limit(self): |
| 374 | limit = self._get_ctx_size_for_history() |
| 375 | total = self.get_tokens() |
| 376 | return total > limit |
| 377 | |
| 378 | def get_bulks_tokens(self) -> int: |
| 379 | return sum(record.get_tokens() for record in self.bulks) |
| 380 | |
| 381 | def get_topics_tokens(self) -> int: |
| 382 | return sum(record.get_tokens() for record in self.topics) |
| 383 | |
| 384 | def get_current_topic_tokens(self) -> int: |
| 385 | return self.current.get_tokens() |
| 386 | |
| 387 | def add_message( |
| 388 | self, |
| 389 | ai: bool, |
| 390 | content: MessageContent, |
| 391 | tokens: int = 0, |
| 392 | id: str = "", |
| 393 | metadata: dict[str, Any] | None = None, |
| 394 | ) -> Message: |
| 395 | self.counter += 1 |
| 396 | return self.current.add_message( |
| 397 | ai, |
| 398 | content=content, |
| 399 | tokens=tokens, |
| 400 | id=id, |
| 401 | metadata=metadata, |
| 402 | sequence=self.counter, |
| 403 | ) |
| 404 | |
| 405 | def new_topic(self): |
| 406 | if self.current.messages: |
| 407 | self.topics.append(self.current) |
| 408 | self.current = Topic(history=self) |
| 409 | |
| 410 | def output(self) -> list[OutputMessage]: |
| 411 | self.trim_embeds(self._get_max_embeds()) |
| 412 | result: list[OutputMessage] = [] |
| 413 | result += [m for b in self.bulks for m in b.output()] |
| 414 | result += [m for t in self.topics for m in t.output()] |
| 415 | result += self.current.output() |
| 416 | return result |
| 417 | |
| 418 | def messages_since(self, sequence: int) -> list[Message]: |
| 419 | return [ |
| 420 | message |
| 421 | for message in self.all_messages() |
| 422 | if int(message.sequence or 0) > int(sequence or 0) |
| 423 | ] |
| 424 | |
| 425 | def all_messages(self) -> list[Message]: |
| 426 | messages: list[Message] = [] |
| 427 | for bulk in self.bulks: |
| 428 | messages.extend(_messages_from_record(bulk)) |
| 429 | for topic in self.topics: |
| 430 | messages.extend(topic.messages) |
| 431 | messages.extend(self.current.messages) |
| 432 | return messages |
| 433 | |
| 434 | def latest_llm_result_for_model(self, provider_model_key: str): |
| 435 | from helpers.llm_result import result_from_metadata |
| 436 | |
| 437 | for message in reversed(self.all_messages()): |
| 438 | if not message.ai: |
| 439 | continue |
| 440 | result = result_from_metadata(message.metadata) |
| 441 | if not result: |
| 442 | continue |
| 443 | if result.provider_model_key == provider_model_key and result.response_id: |
| 444 | return result |
| 445 | return None |
| 446 | |
| 447 | def trim_embeds(self, max_embeds: int) -> int: |
| 448 | if max_embeds == -1: |
| 449 | return 0 |
| 450 | |
| 451 | embeds_count = 0 |
| 452 | removed = 0 |
| 453 | |
| 454 | for record in reversed(self.bulks + self.topics + [self.current]): |
| 455 | embeds_count, removed_now = self._trim_embeds_in_record(record, embeds_count, max_embeds) |
| 456 | removed += removed_now |
| 457 | |
| 458 | return removed |
| 459 | |
| 460 | def remove_all_embeds(self) -> int: |
| 461 | return self.trim_embeds(0) |
| 462 | |
| 463 | def _trim_embeds_in_record( |
| 464 | self, record: Record, embeds_count: int, max_embeds: int |
| 465 | ) -> tuple[int, int]: |
| 466 | if isinstance(record, Message): |
| 467 | if record.summary: |
| 468 | return embeds_count, 0 |
| 469 | |
| 470 | if not _is_raw_message(record.content): |
| 471 | return embeds_count, 0 |
| 472 | |
| 473 | raw_message = cast(dict[str, Any], record.content) |
| 474 | raw_content = raw_message.get("raw_content", []) |
| 475 | if not isinstance(raw_content, list): |
| 476 | return embeds_count, 0 |
| 477 | |
| 478 | embeds_in_message = sum(1 for item in raw_content if _is_embedded_data(item)) |
| 479 | if embeds_in_message <= 0: |
| 480 | return embeds_count, 0 |
| 481 | |
| 482 | if embeds_count + embeds_in_message > max_embeds: |
| 483 | record.set_summary("embedded data removed") |
| 484 | return embeds_count + embeds_in_message, embeds_in_message |
| 485 | |
| 486 | return embeds_count + embeds_in_message, 0 |
| 487 | |
| 488 | if isinstance(record, Topic): |
| 489 | removed = 0 |
| 490 | for message in reversed(record.messages): |
| 491 | embeds_count, removed_now = self._trim_embeds_in_record(message, embeds_count, max_embeds) |
| 492 | removed += removed_now |
| 493 | return embeds_count, removed |
| 494 | |
| 495 | if isinstance(record, Bulk): |
| 496 | removed = 0 |
| 497 | for nested in reversed(record.records): |
| 498 | embeds_count, removed_now = self._trim_embeds_in_record(nested, embeds_count, max_embeds) |
| 499 | removed += removed_now |
| 500 | return embeds_count, removed |
| 501 | |
| 502 | return embeds_count, 0 |
| 503 | |
| 504 | @staticmethod |
| 505 | def from_dict(data: dict, history: "History"): |
| 506 | history.counter = data.get("counter", 0) |
| 507 | history.bulks = [Bulk.from_dict(b, history=history) for b in data["bulks"]] |
| 508 | history.topics = [Topic.from_dict(t, history=history) for t in data["topics"]] |
| 509 | history.current = Topic.from_dict(data["current"], history=history) |
| 510 | return history |
| 511 | |
| 512 | def to_dict(self): |
| 513 | return { |
| 514 | "_cls": "History", |
| 515 | "counter": self.counter, |
| 516 | "bulks": [b.to_dict() for b in self.bulks], |
| 517 | "topics": [t.to_dict() for t in self.topics], |
| 518 | "current": self.current.to_dict(), |
| 519 | } |
| 520 | |
| 521 | def serialize(self): |
| 522 | data = self.to_dict() |
| 523 | return _json_dumps(data) |
| 524 | |
| 525 | async def compress(self): |
| 526 | compressed = False |
| 527 | total = self._get_ctx_size_for_history() |
| 528 | curr, hist, bulk = ( |
| 529 | self.get_current_topic_tokens(), |
| 530 | self.get_topics_tokens(), |
| 531 | self.get_bulks_tokens(), |
| 532 | ) |
| 533 | if (curr + hist + bulk) <= total: |
| 534 | return False |
| 535 | |
| 536 | target = total * COMPRESSION_TARGET_RATIO |
| 537 | prev_total = curr + hist + bulk + 1 |
| 538 | while True: |
| 539 | curr, hist, bulk = ( |
| 540 | self.get_current_topic_tokens(), |
| 541 | self.get_topics_tokens(), |
| 542 | self.get_bulks_tokens(), |
| 543 | ) |
| 544 | |
| 545 | # safeguard against infinite loop in case LLM bloats the summary for some reason |
| 546 | if (curr + hist + bulk) >= prev_total: |
| 547 | break |
| 548 | prev_total = curr + hist + bulk |
| 549 | |
| 550 | ratios = [ |
| 551 | (curr, CURRENT_TOPIC_RATIO, "current_topic"), |
| 552 | (hist, HISTORY_TOPIC_RATIO, "history_topic"), |
| 553 | (bulk, HISTORY_BULK_RATIO, "history_bulk"), |
| 554 | ] |
| 555 | ratios = sorted(ratios, key=lambda x: (x[0] / target) / x[1], reverse=True) |
| 556 | compressed_part = False |
| 557 | for ratio in ratios: |
| 558 | if ratio[0] > ratio[1] * target: |
| 559 | over_part = ratio[2] |
| 560 | if over_part == "current_topic": |
| 561 | compressed_part = await self.current.compress() |
| 562 | elif over_part == "history_topic": |
| 563 | compressed_part = await self.compress_topics() |
| 564 | else: |
| 565 | compressed_part = await self.compress_bulks() |
| 566 | if compressed_part: |
| 567 | break |
| 568 | |
| 569 | if compressed_part: |
| 570 | compressed = True |
| 571 | continue |
| 572 | else: |
| 573 | return compressed |
| 574 | return compressed |
| 575 | |
| 576 | async def compress_topics(self) -> bool: |
| 577 | |
| 578 | # 1. first identify large messages and compress them cheaply |
| 579 | for topic in self.topics: |
| 580 | if topic.compress_large_messages(HISTORY_TOPIC_RATIO*LARGE_MESSAGE_TO_HISTORY_TOPIC_RATIO): |
| 581 | return True |
| 582 | |
| 583 | # 2. summarize topics attention window one by one |
| 584 | for topic in self.topics: |
| 585 | if await topic.compress_attention(HISTORY_TOPIC_ATTENTION_COMPRESSION): |
| 586 | return True |
| 587 | |
| 588 | # 3. move oldest topics to bulks in chunks |
| 589 | if self.topics: |
| 590 | count = TOPICS_MERGE_COUNT if len(self.topics) >= TOPICS_MERGE_COUNT else 1 |
| 591 | chunk = self.topics[:count] |
| 592 | bulk = Bulk(history=self) |
| 593 | bulk.records.extend(chunk) |
| 594 | await bulk.summarize() |
| 595 | self.bulks.append(bulk) |
| 596 | self.topics[:count] = [] |
| 597 | return True |
| 598 | return False |
| 599 | |
| 600 | async def compress_bulks(self): |
| 601 | # merge bulks if possible |
| 602 | compressed = await self.merge_bulks_by(BULK_MERGE_COUNT) |
| 603 | # remove oldest bulk if necessary |
| 604 | if not compressed: |
| 605 | self.bulks.pop(0) |
| 606 | return True |
| 607 | return compressed |
| 608 | |
| 609 | async def merge_bulks_by(self, count: int): |
| 610 | # if bulks is empty, return False |
| 611 | if len(self.bulks) == 0: |
| 612 | return False |
| 613 | # merge bulks in groups of count, even if there are fewer than count |
| 614 | bulks = await asyncio.gather( |
| 615 | *[ |
| 616 | self.merge_bulks(self.bulks[i : i + count]) |
| 617 | for i in range(0, len(self.bulks), count) |
| 618 | ] |
| 619 | ) |
| 620 | self.bulks = bulks |
| 621 | return True |
| 622 | |
| 623 | async def merge_bulks(self, bulks: list[Bulk]) -> Bulk: |
| 624 | bulk = Bulk(history=self) |
| 625 | bulk.records = cast(list[Record], bulks) |
| 626 | await bulk.summarize() |
| 627 | return bulk |
| 628 | |
| 629 | def _get_ctx_size_for_history(self) -> int: |
| 630 | chat_cfg = get_chat_model_config(self.agent) |
| 631 | ctx_length = int(chat_cfg.get("ctx_length", 128000)) |
| 632 | ctx_history = float(chat_cfg.get("ctx_history", 0.7)) |
| 633 | return int(ctx_length * ctx_history) |
| 634 | |
| 635 | def _get_max_embeds(self) -> int: |
| 636 | chat_cfg = get_chat_model_config(self.agent) |
| 637 | if not chat_cfg.get("vision", False): |
| 638 | return 0 |
| 639 | |
| 640 | max_embeds = int(chat_cfg.get("max_embeds", 10)) |
| 641 | if max_embeds <= 0: |
| 642 | max_embeds = -1 |
| 643 | return max_embeds |
| 644 | |
| 645 | |
| 646 | |
| 647 | def deserialize_history(json_data: str, agent) -> History: |
| 648 | history = History(agent=agent) |
| 649 | if json_data: |
| 650 | data = _json_loads(json_data) |
| 651 | history = History.from_dict(data, history=history) |
| 652 | return history |
| 653 | |
| 654 | |
| 655 | def _stringify_output(output: OutputMessage, ai_label="ai", human_label="human"): |
| 656 | return f'{ai_label if output["ai"] else human_label}: {_stringify_content(output["content"])}' |
| 657 | |
| 658 | |
| 659 | def _stringify_content(content: MessageContent) -> str: |
| 660 | # already a string |
| 661 | if isinstance(content, str): |
| 662 | return content |
| 663 | |
| 664 | # raw messages return preview or trimmed json |
| 665 | if _is_raw_message(content): |
| 666 | raw_message = cast(dict[str, Any], content) |
| 667 | preview = raw_message.get("preview") |
| 668 | if isinstance(preview, str) and preview: |
| 669 | return preview |
| 670 | text = _json_dumps(content) |
| 671 | if len(text) > RAW_MESSAGE_OUTPUT_TEXT_TRIM: |
| 672 | return text[:RAW_MESSAGE_OUTPUT_TEXT_TRIM] + "... TRIMMED" |
| 673 | return text |
| 674 | |
| 675 | # regular messages of non-string are dumped as json |
| 676 | return _json_dumps(content) |
| 677 | |
| 678 | |
| 679 | def _output_content_langchain(content: MessageContent): |
| 680 | if isinstance(content, str): |
| 681 | return content |
| 682 | if _is_raw_message(content): |
| 683 | raw_content = cast(dict[str, Any], content).get("raw_content") |
| 684 | return raw_content if raw_content is not None else _json_dumps(content) |
| 685 | try: |
| 686 | return _json_dumps(content) |
| 687 | except Exception as e: |
| 688 | raise e |
| 689 | |
| 690 | |
| 691 | def group_outputs_abab(outputs: list[OutputMessage]) -> list[OutputMessage]: |
| 692 | result = [] |
| 693 | for out in outputs: |
| 694 | if result and result[-1]["ai"] == out["ai"]: |
| 695 | result[-1] = OutputMessage( |
| 696 | ai=result[-1]["ai"], |
| 697 | content=_merge_outputs(result[-1]["content"], out["content"]), |
| 698 | ) |
| 699 | else: |
| 700 | result.append(out) |
| 701 | return result |
| 702 | |
| 703 | |
| 704 | def group_messages_abab(messages: list[BaseMessage]) -> list[BaseMessage]: |
| 705 | result = [] |
| 706 | for msg in messages: |
| 707 | if result and isinstance(result[-1], type(msg)): |
| 708 | # create new instance of the same type with merged content |
| 709 | result[-1] = type(result[-1])(content=_merge_outputs(result[-1].content, msg.content)) # type: ignore |
| 710 | else: |
| 711 | result.append(msg) |
| 712 | return result |
| 713 | |
| 714 | |
| 715 | def output_langchain(messages: list[OutputMessage]): |
| 716 | result = [] |
| 717 | for m in messages: |
| 718 | content = _output_content_langchain(content=m["content"]) |
| 719 | if not content or (isinstance(content, str) and not content.strip()): |
| 720 | continue # skip empty messages, models |
| 721 | if m["ai"]: |
| 722 | result.append(AIMessage(content)) # type: ignore |
| 723 | else: |
| 724 | result.append(HumanMessage(content)) # type: ignore |
| 725 | # ensure message type alternation |
| 726 | result = group_messages_abab(result) |
| 727 | while result and isinstance(result[0], AIMessage): |
| 728 | result.pop(0) |
| 729 | return result |
| 730 | |
| 731 | |
| 732 | def output_text(messages: list[OutputMessage], ai_label="ai", human_label="human"): |
| 733 | return "\n".join(_stringify_output(o, ai_label, human_label) for o in messages) |
| 734 | |
| 735 | |
| 736 | def clear_responses_provider_state(agent) -> None: |
| 737 | key = getattr(agent, "DATA_NAME_RESPONSES_STATE", "responses_state") |
| 738 | get_data = getattr(agent, "get_data", None) |
| 739 | set_data = getattr(agent, "set_data", None) |
| 740 | if not callable(get_data) or not callable(set_data): |
| 741 | return |
| 742 | |
| 743 | state = get_data(key) |
| 744 | if not isinstance(state, dict): |
| 745 | return |
| 746 | |
| 747 | state = dict(state) |
| 748 | removed = False |
| 749 | for field in ("response_id", "previous_response_id"): |
| 750 | if field in state: |
| 751 | state.pop(field, None) |
| 752 | removed = True |
| 753 | |
| 754 | if removed: |
| 755 | set_data(key, state) |
| 756 | |
| 757 | |
| 758 | def _merge_outputs(a: MessageContent, b: MessageContent) -> MessageContent: |
| 759 | if isinstance(a, str) and isinstance(b, str): |
| 760 | return a + "\n" + b |
| 761 | |
| 762 | def make_list(obj: MessageContent) -> list[MessageContent]: |
| 763 | if isinstance(obj, list): |
| 764 | return obj # type: ignore |
| 765 | if isinstance(obj, dict): |
| 766 | return [obj] |
| 767 | if isinstance(obj, str): |
| 768 | return [{"type": "text", "text": obj}] |
| 769 | return [obj] |
| 770 | |
| 771 | a = make_list(a) |
| 772 | b = make_list(b) |
| 773 | |
| 774 | return cast(MessageContent, a + b) |
| 775 | |
| 776 | |
| 777 | def _merge_properties( |
| 778 | a: Dict[str, MessageContent], b: Dict[str, MessageContent] |
| 779 | ) -> Dict[str, MessageContent]: |
| 780 | result = a.copy() |
| 781 | for k, v in b.items(): |
| 782 | if k in result: |
| 783 | result[k] = _merge_outputs(result[k], v) |
| 784 | else: |
| 785 | result[k] = v |
| 786 | return result |
| 787 | |
| 788 | |
| 789 | def _is_raw_message(obj: object) -> bool: |
| 790 | return isinstance(obj, Mapping) and "raw_content" in obj |
| 791 | |
| 792 | |
| 793 | def _is_embedded_data(obj: object) -> bool: |
| 794 | return isinstance(obj, Mapping) and obj.get("type") == "image_url" |
| 795 | |
| 796 | |
| 797 | def _messages_from_record(record: Record) -> list[Message]: |
| 798 | if isinstance(record, Message): |
| 799 | return [record] |
| 800 | if isinstance(record, Topic): |
| 801 | return list(record.messages) |
| 802 | if isinstance(record, Bulk): |
| 803 | messages: list[Message] = [] |
| 804 | for nested in record.records: |
| 805 | messages.extend(_messages_from_record(nested)) |
| 806 | return messages |
| 807 | return [] |
| 808 | |
| 809 | |
| 810 | def _json_dumps(obj): |
| 811 | return json.dumps(obj, ensure_ascii=False, separators=(",", ":")) |
| 812 | |
| 813 | |
| 814 | def _json_loads(obj): |
| 815 | return json.loads(obj) |