codekingpro/portable-devtools

1"""2Portable RAG (Retrieval-Augmented Generation) System3=====================================================4A completely portable, offline-capable RAG setup using ChromaDB and FastEmbed.5 6Usage:7  1. Add a Document:8     python rag_system.py ingest my_document.pdf9     python rag_system.py ingest my_notes.txt10 11  2. Query the Knowledge Base:12     python rag_system.py query "What is the main topic?"13"""14 15import sys16import os17import argparse18from pathlib import Path19 20# Fix for ChromaDB SQLite requirement in some environments21import sqlite322import chromadb23 24# LangChain components25from langchain_community.document_loaders import PyPDFLoader, TextLoader26from langchain_text_splitters import RecursiveCharacterTextSplitter27from langchain_community.vectorstores import Chroma28 29# HuggingFace for Embeddings ^& LLM Generation30from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings31from langchain_core.prompts import ChatPromptTemplate32from langchain_core.runnables import RunnablePassthrough33from langchain_core.output_parsers import StrOutputParser34 35# Configuration36DEVTOOLS_ROOT = Path(os.environ.get("DEVTOOLS_ROOT", os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))37RAG_DATA_DIR = DEVTOOLS_ROOT / "rag_data"38RAG_DATA_DIR.mkdir(exist_ok=True)39 40 41# 1. Initialize Embeddings (LOCAL - No Token Needed)42print("[INFO] Loading Embedding Model (Local - This may take a moment on first run)...")43embeddings = HuggingFaceEmbeddings(44    model_name="sentence-transformers/all-MiniLM-L6-v2",45    model_kwargs={'device': 'cpu'}46)47 48# 2. Initialize Chroma Vector Database49print(f"[INFO] Connecting to ChromaDB at {RAG_DATA_DIR}...")50vectorstore = Chroma(51    embedding_function=embeddings,52    persist_directory=str(RAG_DATA_DIR)53)54 55 56def ingest_document(file_path: str):57    """Reads a file, splits it into chunks, and saves to Vector DB."""58    path = Path(file_path)59    if not path.exists():60        print(f"[ERROR] File not found: {file_path}")61        sys.exit(1)62 63    print(f"\n[1/3] Loading document: {path.name}...")64    if path.suffix.lower() == '.pdf':65        loader = PyPDFLoader(str(path))66    elif path.suffix.lower() == '.txt':67        loader = TextLoader(str(path), encoding='utf-8')68    else:69        print("[ERROR] Unsupported file type. Please use .pdf or .txt")70        sys.exit(1)71 72    docs = loader.load()73    print(f"      Loaded {len(docs)} pages/sections.")74 75    print("[2/3] Splitting into chunks...")76    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)77    splits = text_splitter.split_documents(docs)78    print(f"      Created {len(splits)} chunks.")79 80    print("[3/3] Generating vectors and saving to database...")81    vectorstore.add_documents(splits)82    print("\n[SUCCESS] Document ingested successfully!")83 84 85def format_docs(docs):86    return "\n\n".join(doc.page_content for doc in docs)87 88 89def query_rag(query: str):90    """Retrieves relevant chunks and optionally uses an LLM to generate an answer."""91    print(f"\n[Q] {query}\n")92 93    # Step 1: Retrieval94    print("─── RETRIEVED CONTEXT ────────────────────")95    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})96    results = retriever.invoke(query)97    98    if not results:99        print("No relevant context found in the database.")100        return101 102    for i, doc in enumerate(results):103        source = doc.metadata.get('source', 'Unknown')104        page = doc.metadata.get('page', '')105        page_info = f" (Page {page})" if page else ""106        print(f"[{i+1}] Source: {os.path.basename(source)}{page_info}")107        print(f"    {doc.page_content[:300]}...\n")108 109    # Step 2: Generation (Optional - Requires Token)110    hf_token = os.environ.get("HF_TOKEN")111    if hf_token:112        print("\n─── AI GENERATED ANSWER ──────────────────")113        try:114            llm = HuggingFaceEndpoint(115                repo_id="meta-llama/Meta-Llama-3-8B-Instruct",116                task="text-generation",117                max_new_tokens=512,118                huggingfacehub_api_token=hf_token119            )120            121            template = """You are an assistant for question-answering tasks. 122Use the following pieces of retrieved context to answer the question. 123If you don't know the answer, say that you don't know. 124Keep the answer concise and accurate based ONLY on the context.125 126Context: {context}127 128Question: {question}129 130Answer:"""131            prompt = ChatPromptTemplate.from_template(template)132            133            rag_chain = (134                {"context": retriever | format_docs, "question": RunnablePassthrough()}135                | prompt136                | llm137                | StrOutputParser()138            )139            140            response = rag_chain.invoke(query)141            print(response.strip())142            143        except Exception as e:144            print(f"[WARN] Could not generate AI answer. (Are you logged into Hugging Face?)")145            print(f"       Error: {e}")146    else:147        print("\n[NOTE] HuggingFace token not found. Only returning retrieved documents.")148        print("       Run 'hf auth login' in your terminal to enable AI Generation.")149 150 151def main():152    parser = argparse.ArgumentParser(description="Portable RAG System")153    subparsers = parser.add_subparsers(dest="command", help="Commands")154 155    # Ingest command156    ingest_parser = subparsers.add_parser("ingest", help="Ingest a document into the database")157    ingest_parser.add_argument("file", type=str, help="Path to PDF or TXT file")158 159    # Query command160    query_parser = subparsers.add_parser("query", help="Query the database")161    query_parser.add_argument("query", type=str, help="Your question")162 163    args = parser.parse_args()164 165    if args.command == "ingest":166        ingest_document(args.file)167    elif args.command == "query":168        query_rag(args.query)169    else:170        parser.print_help()171 172 173if __name__ == "__main__":174    main()175