Personal Knowledge Base Agent

Copy all # Personal Knowledge Base RAG Agent # Stack: LlamaIndex + Ollama (local, private) + ChromaDB # Deploy: runs entirely on your laptop, zero API cost # Install: pip install llama-index llama-index-llms-ollama llama-index-embeddings-ollama chromadb gradio import os from pathlib import Path from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings from llama_index.core.storage.storage_context import StorageContext from llama_index.llms.ollama import Ollama from llama_index.embeddings.ollama import OllamaEmbedding from llama_index.vector_stores.chroma import ChromaVectorStore import chromadb import gradio as gr # ── CONFIG ──────────────────────────────────────────────────────── DOCS_DIR = "./knowledge_base" # Put all your notes/PDFs here CHROMA_DIR = "./chroma_db" # Where the vector index is stored LLM_MODEL = "llama3.2" # Any model in Ollama EMBED_MODEL = "nomic-embed-text" # Fast local embedding model # ── SETUP LOCAL LLM + EMBEDDINGS ───────────────────────────────── Settings.llm = Ollama(model=LLM_MODEL, request_timeout=120.0) Settings.embed_model = OllamaEmbedding(model_name=EMBED_MODEL) # ── BUILD OR LOAD VECTOR INDEX ──────────────────────────────────── def get_index(): chroma_client = chromadb.PersistentClient(path=CHROMA_DIR) chroma_collection = chroma_client.get_or_create_collection("knowledge_base") vector_store = ChromaVectorStore(chroma_collection=chroma_collection) # If index exists, load it if len(chroma_collection.get()["ids"]) > 0: print(f"Loading existing index with {len(chroma_collection.get()['ids'])} chunks...") storage_context = StorageContext.from_defaults(vector_store=vector_store) return VectorStoreIndex.from_vector_store(vector_store, storage_context=storage_context) # Otherwise, build from documents print(f"Building index from {DOCS_DIR}...") os.makedirs(DOCS_DIR, exist_ok=True) if not any(Path(DOCS_DIR).iterdir()): print("⚠️ No documents found. Add files to ./knowledge_base/ and restart.") return None documents = SimpleDirectoryReader( DOCS_DIR, recursive=True, required_exts=[".pdf", ".txt", ".md", ".docx"] ).load_data() print(f"Loaded {len(documents)} documents. Building index...") storage_context = StorageContext.from_defaults(vector_store=vector_store) index = VectorStoreIndex.from_documents( documents, storage_context=storage_context, show_progress=True ) print("✅ Index built successfully!") return index # ── QUERY ENGINE WITH CITATIONS ─────────────────────────────────── def create_query_engine(index): return index.as_query_engine( similarity_top_k=5, # Retrieve top 5 most relevant chunks streaming=False, response_mode="compact", # Compact answer with citations verbose=True ) # ── CHAT FUNCTION ───────────────────────────────────────────────── index = get_index() query_engine = create_query_engine(index) if index else None def chat(message, history): if not query_engine: return "⚠️ No documents found. Add files to ./knowledge_base/ directory." try: response = query_engine.query(message) # Format response with sources answer = str(response) sources = [] for node in response.source_nodes: filename = node.metadata.get("file_name", "Unknown") page = node.metadata.get("page_label", "") score = round(node.score, 3) if node.score else "N/A" sources.append(f"📄 {filename} {f'(p.{page})' if page else ''} — relevance: {score}") if sources: answer += "\n\n**Sources:**\n" + "\n".join(sources[:3]) return answer except Exception as e: return f"Error: {str(e)}" def reindex(progress=gr.Progress()): """Re-index all documents (run when you add new files)""" import shutil if os.path.exists(CHROMA_DIR): shutil.rmtree(CHROMA_DIR) global index, query_engine index = get_index() query_engine = create_query_engine(index) if index else None return "✅ Re-indexed successfully!" # ── GRADIO UI ───────────────────────────────────────────────────── with gr.Blocks(title="Knowledge Base Agent", theme=gr.themes.Soft()) as demo: gr.Markdown("# 🧠 Personal Knowledge Base Agent") gr.Markdown("Chat with all your notes, PDFs, and documents. Fully local — no API keys needed.") with gr.Row(): with gr.Column(scale=4): chatbot = gr.ChatInterface( chat, examples=[ "Summarise the key ideas in my documents", "What did I note about machine learning?", "Find everything related to Python", "What are my action items from recent notes?" ], title="" ) with gr.Column(scale=1): gr.Markdown("### Controls") reindex_btn = gr.Button("🔄 Re-index docs", variant="secondary") status = gr.Textbox(label="Status", interactive=False) reindex_btn.click(reindex, outputs=status) doc_count = len(list(Path(DOCS_DIR).rglob("*"))) if Path(DOCS_DIR).exists() else 0 gr.Markdown(f"**Documents folder:** ./knowledge_base/\n**Files found:** {doc_count}") if __name__ == "__main__": print("\n🚀 Starting Knowledge Base Agent...") print("📁 Add your documents to: ./knowledge_base/") print("🌐 Opening at: http://localhost:7860") demo.launch(server_name="0.0.0.0", server_port=7860)