Advanced Memory System

from pathlib import Path
import sys

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 (
    AgentMemory,
    ConversationMemory,
    SemanticMemory,
    EntityMemory,
    InMemoryVectorStore,
    Entity,
)
from shipit_agent.models import Message

# Create 30 messages to demonstrate strategies
messages = [
    Message(
        role="user" if i % 2 == 0 else "assistant",
        content=f"Message #{i}: {'Question' if i % 2 == 0 else 'Answer'} about topic {i}",
    )
    for i in range(30)
]

# --- Buffer strategy: keep everything ---
buffer_mem = ConversationMemory(strategy="buffer")
buffer_mem.add_many(messages)
print(f"Buffer: keeps all {len(buffer_mem.get_messages())} messages")

# --- Window strategy: keep last N ---
window_mem = ConversationMemory(strategy="window", window_size=5)
window_mem.add_many(messages)
window_msgs = window_mem.get_messages()
print(
    f"Window(5): keeps {len(window_msgs)} messages — {window_msgs[0].content[:30]}... to {window_msgs[-1].content[:30]}..."
)

# --- Token strategy: keep within budget ---
token_mem = ConversationMemory(strategy="token", max_tokens=100)
token_mem.add_many(messages)
token_msgs = token_mem.get_messages()
print(f"Token(100): keeps {len(token_msgs)} messages within ~100 token budget")

# --- Summary strategy: summarize old messages ---
summary_mem = ConversationMemory(
    strategy="summary", window_size=5
)  # no LLM = placeholder summary
summary_mem.add_many(messages)
summary_msgs = summary_mem.get_messages()
print(
    f"Summary: {len(summary_msgs)} messages (1 summary + {len(summary_msgs)-1} recent)"
)
print(f"  Summary: {summary_msgs[0].content[:80]}...")

import hashlib


# Simple embedding function (hash-based — for demo; use OpenAI/sentence-transformers in production)
def simple_embed(text: str) -> list[float]:
    """Demo embedding: creates a 16-dim vector from text hash. NOT for production."""
    h = hashlib.sha256(text.lower().encode()).digest()
    return [float(b) / 255.0 for b in h[:16]]


# Create semantic memory with in-memory vector store
semantic = SemanticMemory(
    vector_store=InMemoryVectorStore(),
    embedding_fn=simple_embed,
    top_k=3,
)

# Store facts
facts = [
    "Python was created by Guido van Rossum in 1991",
    "JavaScript was originally created for Netscape Navigator in 1995",
    "Rust focuses on memory safety without garbage collection",
    "Go was designed at Google for cloud infrastructure",
    "TypeScript adds static types to JavaScript",
    "SHIPIT Agent is a powerful Python agent framework",
    "FastAPI is a modern Python web framework using type hints",
    "Docker containers package applications with dependencies",
]

for fact in facts:
    semantic.add(fact)

# Search by meaning
print("=== Search: 'Python programming' ===")
for r in semantic.search("Python programming"):
    print(f"  [{r.score:.3f}] {r.text}")

print("\n=== Search: 'web frameworks' ===")
for r in semantic.search("web frameworks"):
    print(f"  [{r.score:.3f}] {r.text}")

entity_mem = EntityMemory()

# Add entities
entity_mem.add(
    Entity(name="Alice", entity_type="person", context="Lead engineer on Project Atlas")
)
entity_mem.add(
    Entity(name="Bob", entity_type="person", context="Data scientist, ML team")
)
entity_mem.add(
    Entity(
        name="Project Atlas",
        entity_type="project",
        context="Kubernetes migration, due Q2 2026",
    )
)
entity_mem.add(
    Entity(
        name="SHIPIT Agent",
        entity_type="tool",
        context="Python agent framework, v1.0.2",
    )
)

# Lookup
alice = entity_mem.get("Alice")
print(f"Alice: {alice.entity_type} — {alice.context}")

# Search
print("\n=== Search 'Kubernetes' ===")
for e in entity_mem.search("Kubernetes"):
    print(f"  {e.name} ({e.entity_type}): {e.context}")

# Updates merge automatically
entity_mem.add(Entity(name="Alice", context="promoted to CTO"))
alice = entity_mem.get("Alice")
print(f"\nAlice after update: {alice.context}")

# List all entities
print(f"\nAll entities: {[e.name for e in entity_mem.all()]}")

# One-line smart defaults
memory = AgentMemory.default(embedding_fn=simple_embed)

# Add conversation messages
memory.add_message(
    Message(role="user", content="I'm working on Project Atlas with Alice")
)
memory.add_message(
    Message(role="assistant", content="Got it! Project Atlas with Alice.")
)

# Add knowledge facts
memory.add_fact("SHIPIT Agent supports 10 LLM providers")
memory.add_fact("Bedrock's gpt-oss-120b is the cheapest reasoning model")

# Add entities
memory.add_entity(Entity(name="Alice", entity_type="person", context="teammate"))

# Query everything
print("Conversation:", len(memory.get_conversation_messages()), "messages")
print(
    "Knowledge search:", [r.text[:50] for r in memory.search_knowledge("LLM providers")]
)
print(
    "Entity lookup:",
    memory.get_entity("Alice").context if memory.get_entity("Alice") else "Not found",
)