File size: 2,438 Bytes
6e02dfb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | 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 |