backend/agents/memory_agent.py

"""
backend/agents/memory_agent.py

The Memory Agent — manages what the system remembers.

Two phases:
  1. RETRIEVE (before task): Pull relevant memories to inform planning
  2. STORE (after task): Extract and store learnings from this execution

Memory types:
  - Episodic: "Last time I searched for X, Y worked well"
  - Semantic: "The capital of France is Paris"
  - Procedural: "To analyze CSV data: use run_python with pandas"
"""
from __future__ import annotations
import asyncio
from datetime import datetime, timezone

from ..state.graph_state import WorkflowState, AgentRole, make_agent_event, make_memory_entry
from ..memory.memory_store import long_term, short_term
from ..core.config import get_settings
from ..core.logger import get_logger

log = get_logger(__name__)


def memory_retrieve_node(state: WorkflowState) -> WorkflowState:
    """Run at start of workflow — retrieve relevant memories."""
    settings = get_settings()
    if not settings.enable_memory:
        return state

    log.info("Memory retrieval", task_id=state["task_id"])

    # Check short-term first (Redis)
    cached_state = short_term.get_state(state["task_id"])
    if cached_state:
        log.info("Restored from short-term cache")
        return {**state, **cached_state, "task_id": state["task_id"]}

    # Use task text directly — no LLM call needed to generate queries
    queries = [state["task"][:80]]

    # Retrieve from long-term memory
    all_memories = []
    for q in queries[:3]:
        mems = asyncio.run(long_term.retrieve(q, limit=3))
        all_memories.extend(mems)

    # Deduplicate
    seen = set()
    unique_memories = []
    for m in all_memories:
        if m["memory_id"] not in seen:
            seen.add(m["memory_id"])
            unique_memories.append(m)

    log.info("Memories retrieved", count=len(unique_memories))

    return {
        **state,
        "memories": unique_memories[:5],
        "events": state["events"] + [
            make_agent_event(
                AgentRole.MEMORY, "memories_retrieved",
                f"Retrieved {len(unique_memories)} relevant memories",
                {"count": len(unique_memories), "queries": queries},
            )
        ],
    }


def memory_store_node(state: WorkflowState) -> WorkflowState:
    """Run at end of workflow — store learnings from this execution."""
    settings = get_settings()
    if not settings.enable_memory:
        return state

    log.info("Memory storage", task_id=state["task_id"])

    # Store task record to DB
    asyncio.run(long_term.store_task(state["task_id"], state["task"], dict(state)))

    # Cache state in Redis for fast resume
    short_term.store_state(state["task_id"], dict(state))

    # Extract and store new memories
    memories_stored = 0
    for mem in state.get("new_memories", []):
        try:
            asyncio.run(long_term.store(mem, task_id=state["task_id"]))
            memories_stored += 1
        except Exception as e:
            log.warning("Memory store failed", error=str(e))

    # No extra LLM call — critic already added memory entries to new_memories

    log.info("Memories stored", count=memories_stored)

    return {
        **state,
        "events": state["events"] + [
            make_agent_event(
                AgentRole.MEMORY, "memories_stored",
                f"Stored {memories_stored} memories from this execution",
                {"count": memories_stored},
            )
        ],
    }