src/librarian.py

import os
import chromadb
from chromadb.utils import embedding_functions

class Librarian:
    def __init__(self, memory_path):
        # We store the vector database in memory/knowledge_db
        self.db_path = os.path.join(memory_path, "knowledge_db")
        
        # Initialize the database client
        self.client = chromadb.PersistentClient(path=self.db_path)
        
        # Create (or get) the collection where we store memories
        # We use the default embedding model (all-MiniLM-L6-v2) which runs locally and is free
        self.collection = self.client.get_or_create_collection(
            name="kael_knowledge",
            metadata={"hnsw:space": "cosine"} # Cosine similarity for better matching
        )

    def add_document(self, filename, markdown_text):
        """Slices a document into chunks and stores them."""
        print(f"  [Librarian] Indexing {filename}...")
        
        # 1. Simple Chunking (Split by double newlines to get paragraphs)
        # In the future, we can use smarter chunkers, but this works great for v1.
        chunks = markdown_text.split("\n\n")
        
        # Filter out tiny chunks (like page numbers or headers)
        valid_chunks = [c for c in chunks if len(c) > 50]
        
        if not valid_chunks:
            return 0

        # 2. Prepare data for ChromaDB
        ids = [f"{filename}_chunk_{i}" for i in range(len(valid_chunks))]
        metadatas = [{"source": filename, "chunk_id": i} for i in range(len(valid_chunks))]
        
        # 3. Add to Database
        # This automatically converts text -> math (vectors)
        self.collection.add(
            documents=valid_chunks,
            ids=ids,
            metadatas=metadatas
        )
        print(f"  [Librarian] Stored {len(valid_chunks)} chunks from {filename}.")
        return len(valid_chunks)

    def query(self, query_text, n_results=3):
        """Searches the database for the most relevant chunks."""
        results = self.collection.query(
            query_texts=[query_text],
            n_results=n_results
        )
        
        # Flatten the results
        retrieved_knowledge = []
        if results['documents']:
            for i, doc in enumerate(results['documents'][0]):
                source = results['metadatas'][0][i]['source']
                retrieved_knowledge.append(f"From {source}: {doc}")
                
        return retrieved_knowledge