Deep Agents with Memory

from pathlib import Path
import sys
import hashlib

ROOT = (
    Path.cwd().resolve().parent
    if Path.cwd().name == "notebooks"
    else Path.cwd().resolve()
)
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from shipit_agent import (
    Agent,
    AgentMemory,
    ConversationMemory,
    SemanticMemory,
    EntityMemory,
    Entity,
    InMemoryVectorStore,
)
from shipit_agent.deep import GoalAgent, Goal, ReflectiveAgent, Supervisor, Worker
from examples.run_multi_tool_agent import build_llm_from_env
from IPython.display import display, Markdown

llm = build_llm_from_env("bedrock")


# Simple embedding function for demo
def embed(text: str) -> list[float]:
    h = hashlib.sha256(text.lower().encode()).digest()
    return [float(b) / 255.0 for b in h[:16]]


print("Setup complete")

# Create a shared memory that persists across runs
memory = AgentMemory(
    conversation=ConversationMemory(strategy="buffer"),
    knowledge=SemanticMemory(vector_store=InMemoryVectorStore(), embedding_fn=embed),
    entities=EntityMemory(),
)

# --- Run 1: First goal ---
print("=== Run 1: First Goal ===\n")
agent1 = GoalAgent(
    llm=llm,
    memory=memory,
    goal=Goal(
        objective="List the top 3 Python web frameworks",
        success_criteria=["Names Django, Flask, and FastAPI"],
        max_steps=2,
    ),
)

result1 = agent1.run()
print(f"Status: {result1.goal_status}")
print(f"Output: {result1.output[:300]}...")

# Store the result as a fact in memory
memory.add_fact(f"Previous research: {result1.output[:500]}")
memory.add_entity(
    Entity(name="Django", entity_type="framework", context="Python web framework")
)
memory.add_entity(
    Entity(
        name="Flask",
        entity_type="framework",
        context="Lightweight Python web framework",
    )
)
memory.add_entity(
    Entity(
        name="FastAPI",
        entity_type="framework",
        context="Modern async Python web framework",
    )
)

print(
    f"\nMemory: {len(memory.get_conversation_messages())} messages, {len(memory.entities.all())} entities"
)

# --- Run 2: Second goal — agent remembers Run 1! ---
print("=== Run 2: Build on Previous Research ===\n")
agent2 = GoalAgent(
    llm=llm,
    memory=memory,  # SAME memory — agent sees previous conversation
    goal=Goal(
        objective="Compare the performance of the frameworks you found earlier",
        success_criteria=[
            "References Django, Flask, and FastAPI",
            "Includes performance comparison",
        ],
        max_steps=2,
    ),
)

result2 = agent2.run()
print(f"Status: {result2.goal_status}")
print(f"Output: {result2.output[:400]}...")

# Check what the agent remembers
print("\n=== Memory State ===")
print(f"Conversation messages: {len(memory.get_conversation_messages())}")
print(f"Entities tracked: {[e.name for e in memory.entities.all()]}")

# Search knowledge
print("\nKnowledge search for 'FastAPI':")
for r in memory.search_knowledge("FastAPI"):
    print(f"  [{r.score:.2f}] {r.text[:80]}...")

# ReflectiveAgent with memory
reflect_memory = AgentMemory(
    conversation=ConversationMemory(
        strategy="summary", summary_llm=llm, window_size=10
    ),
    knowledge=SemanticMemory(vector_store=InMemoryVectorStore(), embedding_fn=embed),
    entities=EntityMemory(),
)

reflective = ReflectiveAgent(
    llm=llm,
    memory=reflect_memory,
    reflection_prompt="Check accuracy and completeness. Score 0-1.",
    max_reflections=2,
    quality_threshold=0.8,
)

# Run 1
print("=== ReflectiveAgent Run 1 ===\n")
result1 = reflective.run("Explain what a REST API is")
print(f"Quality: {result1.final_quality:.2f}")
print(f"Output: {result1.output[:200]}...")

# Store what it learned
reflect_memory.add_fact(f"REST API explanation quality: {result1.final_quality:.2f}")
for r in result1.reflections:
    reflect_memory.add_fact(f"Reflection feedback: {r.feedback[:200]}")

print(
    f"\nMemory: {len(reflect_memory.get_conversation_messages())} msgs, {len(reflect_memory.search_knowledge('REST'))} facts"
)

# Run 2 — agent has context from Run 1
print("\n=== ReflectiveAgent Run 2 (with memory) ===\n")
result2 = reflective.run("Now explain GraphQL and compare it to REST")
print(f"Quality: {result2.final_quality:.2f}")
print(f"Output: {result2.output[:300]}...")

# Shared memory for the entire team
team_memory = AgentMemory(
    conversation=ConversationMemory(strategy="buffer"),
    knowledge=SemanticMemory(vector_store=InMemoryVectorStore(), embedding_fn=embed),
    entities=EntityMemory(),
)

# Pre-load some knowledge
team_memory.add_fact("Q4 2025 revenue: $2.4M, up 15% YoY")
team_memory.add_fact("Customer satisfaction: 92%, churn rate 2.1%")
team_memory.add_fact("Top market: North America (60%), growing in APAC (25%)")
team_memory.add_entity(
    Entity(name="Product Atlas", entity_type="product", context="Main SaaS platform")
)

# Workers share the same memory via history
analyst = Worker(
    name="analyst",
    agent=Agent(
        llm=llm,
        prompt="You are a data analyst. Use the context from previous conversations to provide insights.",
        history=team_memory.get_conversation_messages(),
    ),
)

writer = Worker(
    name="writer",
    agent=Agent(
        llm=llm,
        prompt="You are a report writer. Use data from the analyst's findings.",
        history=team_memory.get_conversation_messages(),
    ),
)

supervisor = Supervisor(llm=llm, workers=[analyst, writer], max_delegations=4)

print("=== Supervisor with Shared Memory ===\n")
print(f"Pre-loaded: {len(team_memory.search_knowledge('revenue'))} facts about revenue")
print(f"Entities: {[e.name for e in team_memory.entities.all()]}\n")

result = supervisor.run("Write a Q4 executive summary using the data we have")

for d in result.delegations:
    print(f"\n[{d.worker}] Round {d.round}")
    print(f"  Task: {d.task[:80]}...")
    print(f"  Output: {d.output[:200]}...")

print("\n=== Final Output ===")
display(Markdown(result.output[:800]))