Memory consolidation

from shipit_agent import Agent
from shipit_agent.memory import (
    AgentMemory, MemoryConsolidator, SemanticMemory, InMemoryVectorStore,
)

# Build memory
memory = AgentMemory(
    knowledge=SemanticMemory(vector_store=InMemoryVectorStore()),
)

# Consolidator — typically a cheap LLM (Haiku, gpt-oss-20b)
consolidator = MemoryConsolidator(
    llm=cheap_llm,
    max_facts_per_pass=8,
    min_messages=6,
    confidence_threshold=0.5,
)

# Build an agent that uses the memory
agent = Agent(llm=opus_llm)

# After a few turns, consolidate
result = consolidator.consolidate(
    memory=memory,
    recent_messages=memory.conversation.get_messages(),
)

for fact in result.facts:
    print(f"[{fact.category}] {fact.text} ({fact.confidence:.2f})")

[preference] User wants concise answers, no preamble
[project]    Auth service uses Argon2 with 12-byte salt
[goal]       Q3 release deadline is Oct 14
[person]     Alice owns the deploy pipeline

# Half-life of 14 days: a 14-day-old fact decays to strength 0.5,
# 28-day-old to 0.25, etc. Facts whose strength falls below
# `forgetting_threshold` (default 0.1) are pruned.

pruned = consolidator.decay(memory.knowledge, half_life_days=14)
print(f"pruned {pruned} stale facts")

strength_new = strength_old × exp(-ln(2) × elapsed_days / half_life_days)

core_facts = consolidator.core_memory(
    memory.knowledge,
    top_k=5,
    min_retrievals=0,   # set ≥1 to require some past retrieval before promoting
)

system_prompt_addition = "\n".join(f"- {f}" for f in core_facts)

agent = Agent(
    llm=opus_llm,
    prompt=BASE_PROMPT + "\n\nKnown facts about this user/project:\n" + system_prompt_addition,
)

results = memory.search_knowledge("Argon2", top_k=3)
consolidator.record_retrieval(
    memory.knowledge,
    [r.text for r in results],
)

class MemoryEnabledAgent:
    def __init__(self, llm, cheap_llm):
        self.memory = AgentMemory.default(llm=llm)
        self.consolidator = MemoryConsolidator(llm=cheap_llm)
        self.turn_count = 0

    def chat(self, user_message: str) -> str:
        # 1. Inject core facts into prompt
        core = self.consolidator.core_memory(self.memory.knowledge, top_k=5)
        agent = Agent(
            llm=opus_llm,
            prompt=BASE_PROMPT + "\n\nKnown facts:\n" + "\n".join(f"- {f}" for f in core),
        )

        # 2. Run the turn
        result = agent.run(user_message)

        # 3. Update conversation memory
        from shipit_agent.models import Message
        self.memory.add_message(Message(role="user", content=user_message))
        self.memory.add_message(Message(role="assistant", content=result.output))
        self.turn_count += 1

        # 4. Every 10 turns, consolidate
        if self.turn_count % 10 == 0:
            self.consolidator.consolidate(
                memory=self.memory,
                recent_messages=self.memory.get_conversation_messages(),
            )

        # 5. Once per day in your real app: decay + prune
        # self.consolidator.decay(self.memory.knowledge, half_life_days=14)

        return result.output

MemoryConsolidator(
    llm=cheap_llm,                # any object with .complete(messages=...)
    max_facts_per_pass=8,         # cap on facts written per consolidation
    min_messages=6,               # don't consolidate short chats
    confidence_threshold=0.5,     # drop low-confidence facts
    forgetting_threshold=0.1,     # below this strength → pruned in decay()
)

MemoryConsolidator(
    *, llm,
    max_facts_per_pass: int = 8,
    min_messages: int = 6,
    confidence_threshold: float = 0.5,
    forgetting_threshold: float = 0.1,
)