Agent Memory & Planning

from dataclasses import dataclass, field

@dataclass
class Memory:
    short_term: list[dict] = field(default_factory=list)
    user_id: str = ""

    def add_message(self, role: str, content: str):
        self.short_term.append({"role": role, "content": content})
        # Summarise if history too long
        if len(self.short_term) > 20:
            self._summarise()

    def _summarise(self):
        """Compress old messages into a summary to free context budget."""
        old = self.short_term[:15]
        recent = self.short_term[15:]
        summary = llm_summarise(old)
        self.short_term = [{"role": "system", "content": f"Conversation summary: {summary}"}] + recent

    def recall_long_term(self, query: str) -> list[dict]:
        """Retrieve relevant long-term memories for this user."""
        qemb = embed(query)
        return vector_store.search(
            embedding=qemb,
            filter={"user_id": self.user_id},
            k=3,
        )

def plan(task: str) -> list[str]:
    """Decompose a complex task into steps."""
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "Decompose the task into 3-7 concrete steps. Return as JSON array of strings."},
            {"role": "user", "content": task},
        ],
        response_format={"type": "json_object"},
    )
    import json
    return json.loads(response.choices[0].message.content)["steps"]

def execute_with_plan(task: str, memory: Memory, tools: dict) -> str:
    """Plan, then execute each step, accumulating results."""
    steps = plan(task)
    memory.add_message("system", f"Plan: {steps}")
    results = []
    for i, step in enumerate(steps):
        memory.add_message("user", f"Step {i+1}: {step}")
        result = run_agent(step, tools, memory)
        results.append(result)
        memory.add_message("assistant", result)
    # Synthesise
    return run_agent(f"Synthesise final answer from: {results}", tools, memory)