62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
"""
|
|
Lab 2.9: Context-Aware Agent -- SOLUTION
|
|
|
|
Implements sliding window + summarization for long agent sessions.
|
|
"""
|
|
|
|
|
|
class ContextManager:
|
|
def __init__(self, max_recent_turns: int = 5, system_prompt: str = ""):
|
|
self.system_prompt = system_prompt
|
|
self.max_recent_turns = max_recent_turns
|
|
self.summary = ""
|
|
self.recent_messages = []
|
|
|
|
def add_message(self, role: str, content: str):
|
|
self.recent_messages.append({"role": role, "content": content})
|
|
|
|
max_messages = self.max_recent_turns * 2
|
|
if len(self.recent_messages) > max_messages:
|
|
oldest = self.recent_messages[:2]
|
|
self.recent_messages = self.recent_messages[2:]
|
|
|
|
if self.summary:
|
|
self.summary += " | "
|
|
self.summary += f"[{oldest[0]['role']}]: {oldest[0]['content'][:100]}..."
|
|
if len(oldest) > 1:
|
|
self.summary += f" -> [{oldest[1]['role']}]: {oldest[1]['content'][:100]}..."
|
|
|
|
def build_context(self) -> list[dict]:
|
|
context = []
|
|
if self.system_prompt:
|
|
context.append({"role": "system", "content": self.system_prompt})
|
|
if self.summary:
|
|
context.append({"role": "system", "content": f"Session summary (earlier context): {self.summary}"})
|
|
context.extend(self.recent_messages)
|
|
return context
|
|
|
|
def summarize(self, messages: list[dict]) -> str:
|
|
parts = []
|
|
for m in messages:
|
|
role = m.get("role", "unknown")
|
|
content = m.get("content", "")[:80]
|
|
parts.append(f"[{role}]: {content}")
|
|
return " | ".join(parts)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cm = ContextManager(system_prompt="You are a helpful assistant.")
|
|
|
|
for i in range(20):
|
|
user_msg = f"User message {i+1} asking about topic {'A B C D'[i % 4]}"
|
|
asst_msg = f"Response {i+1}: Here is information about topic {'A B C D'[i % 4]}"
|
|
cm.add_message("user", user_msg)
|
|
cm.add_message("assistant", asst_msg)
|
|
|
|
context = cm.build_context()
|
|
print(f"Context messages: {len(context)}")
|
|
for m in context:
|
|
role = m["role"]
|
|
content = m["content"][:80]
|
|
print(f" [{role}]: {content}...")
|