rag.py

class RagManager:
    def __init__(self, db_path: str = "./chroma_db"):
        self.db_path = db_path
        
        # Ensure directory exists
        os.makedirs(db_path, exist_ok=True)
        
        # Persistent client stores data to disk
        self.client = chromadb.PersistentClient(path=db_path)
        
        # Use sentence-transformers for local embeddings
        model_name = "all-MiniLM-L6-v2"
        self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=model_name
        )

@property
def collection(self) -> Collection:
    if self._collection is None:
        self._collection = self.client.get_or_create_collection(
            name="world_lore",
            embedding_function=self.embedding_function,
            metadata={"hnsw:space": "cosine"}
        )
    return self._collection

@property
def memory_collection(self) -> Collection:
    if self._memory_collection is None:
        self._memory_collection = self.client.get_or_create_collection(
            name="session_memory",
            embedding_function=self.embedding_function,
            metadata={"hnsw:space": "cosine"}
        )
    return self._memory_collection

def add_lore(
    self, world_id: str, lore_id: str, text: str, metadata: Dict[str, Any] = None
)

rag_manager.add_lore(
    world_id="fantasy",
    lore_id="ruins_001",
    text="The ancient ruins of Kal'dara hold secrets from the First Age. Only the bravest adventurers dare explore their cursed halls.",
    metadata={"category": "locations", "importance": "high"}
)

meta = metadata or {}
meta["world_id"] = world_id  # Ensure world filtering works

self.collection.upsert(
    ids=[f"{world_id}_{lore_id}"],
    documents=[text],
    metadatas=[meta]
)

def query_lore(
    self, world_id: str, query: str, n_results: int = 3, max_chars: int = 1000
) -> Tuple[List[str], bool]

lore_chunks, warning = rag_manager.query_lore(
    world_id="fantasy",
    query="Tell me about the ancient ruins",
    n_results=2
)

if warning:
    print("Warning: Retrieved context is large")

for chunk in lore_chunks:
    print(f"Lore: {chunk}")

results = self.collection.query(
    query_texts=[query],
    n_results=n_results,
    where={"world_id": world_id}  # Filter by world ID
)

if results["documents"] and results["documents"][0]:
    docs = results["documents"][0]
    
    # Check for context window bloat
    total_chars = sum(len(d) for d in docs)
    context_warning = total_chars > max_chars
    
    return docs, context_warning
return [], False

def add_memory(self, session_id: str, memory_id: str, text: str)

import uuid

memory_key = f"elara_a1b2c3d4e5f6"
rag_manager.add_memory(
    session_id=memory_key,
    memory_id=str(uuid.uuid4()),
    text="User: What happened in the forest?\nAI: You discovered ancient markings on the trees, suggesting this place was once sacred."
)

self.memory_collection.upsert(
    ids=[f"{session_id}_{memory_id}"],
    documents=[text],
    metadatas=[{"session_id": session_id}]
)

def query_memory(
    self, session_id: str, query: str, n_results: int = 3, max_chars: int = 1500
) -> Tuple[List[str], bool]

memory_key = f"elara_a1b2c3d4e5f6"
memories, warning = rag_manager.query_memory(
    session_id=memory_key,
    query="What did we discuss about magic?",
    n_results=3
)

for memory in memories:
    print(f"Memory: {memory}")

results = self.memory_collection.query(
    query_texts=[query],
    n_results=n_results,
    where={"session_id": session_id}
)

if results["documents"] and results["documents"][0]:
    docs = results["documents"][0]
    total_chars = sum(len(d) for d in docs)
    context_warning = total_chars > max_chars
    return docs, context_warning
return [], False

def delete_session_memory(self, session_id: str)

memory_key = f"elara_a1b2c3d4e5f6"
rag_manager.delete_session_memory(memory_key)

self.memory_collection.delete(where={"session_id": session_id})

rag_manager = RagManager()

from engine.rag import rag_manager

# Use directly
rag_manager.add_lore("fantasy", "lore_001", "Ancient lore text...")

for data in load_yaml_assets("assets/worlds/*.yaml"):
    w = World(...)
    world_manager.add_world(w)
    
    # Split large lore into chunks
    if w.lore:
        chunks = []
        current_chunk = ""
        chunk_size = 800
        
        for line in w.lore.split('\n'):
            if len(current_chunk) + len(line) > chunk_size and current_chunk:
                chunks.append(current_chunk.strip())
                current_chunk = line + '\n'
            else:
                current_chunk += line + '\n'
        
        if current_chunk.strip():
            chunks.append(current_chunk.strip())
        
        # Add each chunk separately
        for i, chunk in enumerate(chunks):
            rag_manager.add_lore(w.id, f"base_lore_{i}", chunk)

# RAG retrieval
lore_list, _ = rag_manager.query_lore(
    state.ACTIVE_WORLD_ID, prompt, n_results=2
)
mem_key = f"{state.ACTIVE_CHARACTER_ID}_{state.ACTIVE_SESSION_ID}"
mem_list, _ = rag_manager.query_memory(mem_key, prompt, n_results=3)

# Log retrieved chunks
if lore_list:
    log.info(f"\n=== RETRIEVED LORE ({len(lore_list)} chunks) ===")
    for i, chunk in enumerate(lore_list):
        log.info(f"[LORE {i+1}]\n{chunk}\n")

if mem_list:
    log.info(f"\n=== RETRIEVED MEMORY ({len(mem_list)} chunks) ===")
    for i, chunk in enumerate(mem_list):
        log.info(f"[MEMORY {i+1}]\n{chunk}\n")

# Build context
full_context = (
    f"--- RECENT MEMORY ---\n{chr(10).join(mem_list)}" if mem_list else ""
)

if full_content:
    memory_key = f"{char_id}_{session_id}" if session_id else char_id
    rag_manager.add_memory(
        memory_key, str(uuid.uuid4()), f"User: {prompt}\nAI: {full_content}"
    )

import uuid
from engine.rag import rag_manager

# === LORE MANAGEMENT ===

# Add world lore
rag_manager.add_lore(
    world_id="fantasy",
    lore_id="location_001",
    text="The Tower of Echoes stands at the edge of reality, where the veil between worlds is thinnest.",
    metadata={"type": "location", "region": "borderlands"}
)

rag_manager.add_lore(
    world_id="fantasy",
    lore_id="history_001",
    text="During the War of Shadows, the ancient mages sealed away forbidden knowledge in hidden vaults.",
    metadata={"type": "history", "era": "ancient"}
)

# Query lore
query = "Tell me about magical towers"
lore_results, warning = rag_manager.query_lore(
    world_id="fantasy",
    query=query,
    n_results=2
)

print(f"Query: {query}")
print(f"Context Warning: {warning}")
for i, chunk in enumerate(lore_results, 1):
    print(f"\nLore {i}: {chunk}")

# === MEMORY MANAGEMENT ===

# Session setup
char_id = "elara"
session_id = "a1b2c3d4e5f6"
memory_key = f"{char_id}_{session_id}"

# Store conversation exchanges
rag_manager.add_memory(
    session_id=memory_key,
    memory_id=str(uuid.uuid4()),
    text="User: What's your background?\nAI: I'm a wandering scholar who studies ancient magic and forgotten histories."
)

rag_manager.add_memory(
    session_id=memory_key,
    memory_id=str(uuid.uuid4()),
    text="User: Have you been to the Tower of Echoes?\nAI: Yes, I've visited it once. The experience was... unsettling."
)

# Query memory
query = "Have we talked about your travels?"
memory_results, warning = rag_manager.query_memory(
    session_id=memory_key,
    query=query,
    n_results=3
)

print(f"\nMemory Query: {query}")
for i, memory in enumerate(memory_results, 1):
    print(f"\nMemory {i}: {memory}")

# === CLEANUP ===

# Delete all session memory when done
rag_manager.delete_session_memory(memory_key)
print(f"\nDeleted all memory for session: {memory_key}")

self.client = chromadb.PersistentClient(path=db_path)

chroma_db/
├── chroma.sqlite3       # Metadata and document storage
└── [collection_id]/     # Vector embeddings
    └── data_level0.bin

metadata={"hnsw:space": "cosine"}

chunk_size = 800  # characters per chunk

max_chars = 1000  # for lore
max_chars = 1500  # for memory

try:
    results = self.collection.query(...)
    # ...
except Exception as e:
    log.error(f"Error querying lore: {e}")
    return [], False