|
|
from typing import Callable, List, Set, Tuple |
|
|
|
|
|
from injector import inject |
|
|
|
|
|
from taskweaver.config.module_config import ModuleConfig |
|
|
from taskweaver.llm import LLMApi |
|
|
from taskweaver.llm.util import format_chat_message |
|
|
from taskweaver.logging import TelemetryLogger |
|
|
from taskweaver.memory import Round |
|
|
|
|
|
|
|
|
class RoundCompressorConfig(ModuleConfig): |
|
|
def _configure(self) -> None: |
|
|
self._set_name("round_compressor") |
|
|
self.rounds_to_compress = self._get_int("rounds_to_compress", 2) |
|
|
self.rounds_to_retain = self._get_int("rounds_to_retain", 3) |
|
|
|
|
|
assert self.rounds_to_compress > 0, "rounds_to_compress must be greater than 0" |
|
|
assert self.rounds_to_retain > 0, "rounds_to_retain must be greater than 0" |
|
|
|
|
|
|
|
|
class RoundCompressor: |
|
|
@inject |
|
|
def __init__( |
|
|
self, |
|
|
llm_api: LLMApi, |
|
|
config: RoundCompressorConfig, |
|
|
logger: TelemetryLogger, |
|
|
): |
|
|
self.config = config |
|
|
self.processed_rounds: Set[str] = set() |
|
|
self.rounds_to_compress = self.config.rounds_to_compress |
|
|
self.rounds_to_retain = self.config.rounds_to_retain |
|
|
self.previous_summary: str = "None" |
|
|
self.llm_api = llm_api |
|
|
self.logger = logger |
|
|
|
|
|
def compress_rounds( |
|
|
self, |
|
|
rounds: List[Round], |
|
|
rounds_formatter: Callable, |
|
|
use_back_up_engine: bool = False, |
|
|
prompt_template: str = "{PREVIOUS_SUMMARY}, please compress the following rounds", |
|
|
) -> Tuple[str, List[Round]]: |
|
|
remaining_rounds = len(rounds) |
|
|
for _round in rounds: |
|
|
if _round.id in self.processed_rounds: |
|
|
remaining_rounds -= 1 |
|
|
continue |
|
|
break |
|
|
|
|
|
|
|
|
if remaining_rounds < (self.rounds_to_compress + self.rounds_to_retain): |
|
|
return self.previous_summary, rounds[-remaining_rounds:] |
|
|
|
|
|
chat_summary = self._summarize( |
|
|
rounds[-remaining_rounds : -self.rounds_to_retain], |
|
|
rounds_formatter, |
|
|
use_back_up_engine=use_back_up_engine, |
|
|
prompt_template=prompt_template, |
|
|
) |
|
|
|
|
|
if len(chat_summary) > 0: |
|
|
self.previous_summary = chat_summary |
|
|
return chat_summary, rounds[-self.rounds_to_retain :] |
|
|
else: |
|
|
return self.previous_summary, rounds[-remaining_rounds:] |
|
|
|
|
|
def _summarize( |
|
|
self, |
|
|
rounds: List[Round], |
|
|
rounds_formatter: Callable, |
|
|
use_back_up_engine: bool = False, |
|
|
prompt_template: str = "{PREVIOUS_SUMMARY}, please compress the following rounds", |
|
|
) -> str: |
|
|
assert "{PREVIOUS_SUMMARY}" in prompt_template, "Prompt template must contain {PREVIOUS_SUMMARY}" |
|
|
try: |
|
|
chat_history_str = rounds_formatter(rounds) |
|
|
system_instruction = prompt_template.format( |
|
|
PREVIOUS_SUMMARY=self.previous_summary, |
|
|
) |
|
|
prompt = [ |
|
|
format_chat_message("system", system_instruction), |
|
|
format_chat_message("user", chat_history_str), |
|
|
] |
|
|
new_summary = self.llm_api.chat_completion(prompt, use_backup_engine=use_back_up_engine)["content"] |
|
|
self.processed_rounds.update([_round.id for _round in rounds]) |
|
|
return new_summary |
|
|
except Exception as e: |
|
|
self.logger.warning(f"Failed to compress rounds: {e}") |
|
|
return "" |
|
|
|
