| """ |
| BeastBullet + RAG: Personal Knowledge Assistant |
| Query your documents, emails, notes - 100% local |
| """ |
|
|
| import os |
| from pathlib import Path |
| from typing import List, Dict |
| import json |
|
|
| class DocumentRAG: |
| """Document retrieval for personal knowledge.""" |
| |
| def __init__(self, vector_db="chromadb", collection_name="my_knowledge"): |
| self.collection_name = collection_name |
| self.db = self._init_vector_db(vector_db) |
| |
| def _init_vector_db(self, db_type: str): |
| """Initialize vector database.""" |
| if db_type == "chromadb": |
| try: |
| import chromadb |
| client = chromadb.Client() |
| collection = client.get_or_create_collection( |
| name=self.collection_name |
| ) |
| print(f"✅ ChromaDB initialized: {self.collection_name}") |
| return collection |
| except ImportError: |
| print("⚠️ Install chromadb: pip install chromadb") |
| return None |
| else: |
| raise ValueError(f"Unsupported vector DB: {db_type}") |
| |
| def index_directory(self, directory: str, extensions: List[str] = None): |
| """Index all documents in a directory.""" |
| if not self.db: |
| return |
| |
| if extensions is None: |
| extensions = ['.txt', '.md', '.pdf', '.docx', '.py', '.js', '.json'] |
| |
| directory = Path(directory).expanduser() |
| print(f"\n📁 Indexing: {directory}") |
| |
| indexed = 0 |
| for ext in extensions: |
| for file_path in directory.rglob(f"*{ext}"): |
| try: |
| |
| if ext == '.pdf': |
| content = self._read_pdf(file_path) |
| else: |
| with open(file_path, 'r', encoding='utf-8', errors='ignore') as f: |
| content = f.read() |
| |
| if not content.strip(): |
| continue |
| |
| |
| chunks = self._chunk_text(content) |
| |
| |
| for i, chunk in enumerate(chunks): |
| self.db.add( |
| documents=[chunk], |
| ids=[f"{file_path}_{i}"], |
| metadatas=[{ |
| "source": str(file_path), |
| "chunk": i, |
| "total_chunks": len(chunks) |
| }] |
| ) |
| |
| indexed += 1 |
| print(f" ✅ {file_path.name} ({len(chunks)} chunks)") |
| |
| except Exception as e: |
| print(f" ⚠️ Error indexing {file_path.name}: {e}") |
| |
| print(f"\n✅ Indexed {indexed} documents") |
| |
| def _chunk_text(self, text: str, chunk_size: int = 500) -> List[str]: |
| """Split text into chunks.""" |
| words = text.split() |
| chunks = [] |
| |
| for i in range(0, len(words), chunk_size): |
| chunk = ' '.join(words[i:i + chunk_size]) |
| if chunk.strip(): |
| chunks.append(chunk) |
| |
| return chunks |
| |
| def _read_pdf(self, file_path: Path) -> str: |
| """Read PDF file.""" |
| try: |
| import PyPDF2 |
| with open(file_path, 'rb') as f: |
| reader = PyPDF2.PdfReader(f) |
| text = "" |
| for page in reader.pages: |
| text += page.extract_text() |
| return text |
| except ImportError: |
| print("⚠️ Install PyPDF2: pip install PyPDF2") |
| return "" |
| |
| def retrieve(self, query: str, n_results: int = 5) -> Dict: |
| """Retrieve relevant documents.""" |
| if not self.db: |
| return {"documents": [], "sources": []} |
| |
| results = self.db.query( |
| query_texts=[query], |
| n_results=n_results |
| ) |
| |
| |
| documents = results['documents'][0] if results['documents'] else [] |
| metadatas = results['metadatas'][0] if results['metadatas'] else [] |
| |
| sources = [] |
| for meta in metadatas: |
| if meta['source'] not in [s['file'] for s in sources]: |
| sources.append({ |
| "file": meta['source'], |
| "chunks": 1 |
| }) |
| else: |
| for s in sources: |
| if s['file'] == meta['source']: |
| s['chunks'] += 1 |
| |
| return { |
| "documents": documents, |
| "sources": sources, |
| "context": "\n\n".join(documents) |
| } |
|
|
|
|
| class PersonalKnowledgeAssistant: |
| """Personal knowledge assistant with BeastBullet + RAG.""" |
| |
| def __init__(self): |
| from code.integration_example import TinyLlamaIntegration, ExpertLoader |
| |
| print("🧠 Loading BeastBullet...") |
| loader = ExpertLoader("experts/") |
| loader.load_encoder() |
| loader.load_all_experts() |
| |
| self.bb = TinyLlamaIntegration( |
| experts=loader.experts, |
| encoder=loader.encoder |
| ) |
| |
| self.rag = DocumentRAG() |
| print("✅ Personal Knowledge Assistant ready!\n") |
| |
| def index_my_knowledge(self): |
| """Index common personal knowledge locations.""" |
| print("📚 Indexing your knowledge base...") |
| |
| |
| locations = [ |
| "~/Documents", |
| "~/Downloads", |
| "~/Desktop", |
| "~/Notes", |
| ] |
| |
| for location in locations: |
| path = Path(location).expanduser() |
| if path.exists(): |
| self.rag.index_directory(path) |
| |
| def query(self, question: str) -> Dict: |
| """Query your personal knowledge.""" |
| print("=" * 70) |
| print(f"Question: {question}") |
| print("=" * 70) |
| |
| |
| retrieval = self.rag.retrieve(question) |
| |
| if not retrieval["documents"]: |
| print("\n⚠️ No relevant documents found in your knowledge base") |
| result = self.bb.query(question) |
| return result |
| |
| |
| context_query = f"""Based on my personal documents, answer: {question} |
| |
| My Documents: |
| {retrieval['context']} |
| |
| Provide a clear answer based on my documents.""" |
| |
| |
| result = self.bb.query(context_query) |
| result["sources"] = retrieval["sources"] |
| |
| return result |
| |
| def display_result(self, result: Dict): |
| """Display result.""" |
| print("\n" + "=" * 70) |
| print("🧠 Answer:") |
| print("=" * 70) |
| print(result.get("answer", "No answer generated")) |
| |
| if "sources" in result and result["sources"]: |
| print("\n" + "=" * 70) |
| print("📁 Sources from your documents:") |
| print("=" * 70) |
| for source in result["sources"]: |
| print(f" • {source['file']} ({source['chunks']} relevant sections)") |
| |
| print("\n" + "=" * 70) |
| print(f"Experts: {', '.join(result.get('experts_used', []))}") |
| print(f"Confidence: {result.get('confidence', 0):.1%}") |
| print("=" * 70) |
|
|
|
|
| |
| |
| |
|
|
| def setup(): |
| """Setup and index documents.""" |
| assistant = PersonalKnowledgeAssistant() |
| assistant.index_my_knowledge() |
| return assistant |
|
|
|
|
| def interactive(): |
| """Interactive mode.""" |
| assistant = setup() |
| |
| print("\n🧠 Personal Knowledge Assistant - Interactive Mode") |
| print("=" * 70) |
| print("Ask questions about your documents. Type 'exit' to quit.") |
| print("=" * 70) |
| print() |
| |
| while True: |
| try: |
| question = input("You: ").strip() |
| |
| if question.lower() in ['exit', 'quit', 'q']: |
| print("\n👋 Goodbye!") |
| break |
| |
| if not question: |
| continue |
| |
| result = assistant.query(question) |
| assistant.display_result(result) |
| print("\n") |
| |
| except KeyboardInterrupt: |
| print("\n\n👋 Goodbye!") |
| break |
|
|
|
|
| if __name__ == "__main__": |
| interactive() |
|
|
