process_docs.py

#!/usr/bin/env python3
"""
Script to process documents and create vector store for the RAG application.
This is a Python version of the ingest_and_build.ipynb notebook.
"""

import os
import sys
from pathlib import Path
import json

# Add the parent directory to the path so we can import our modules
sys.path.append('..')

# Import our modules
try:
    from app.retriever import LocalRetriever
    from app.utils import load_documents_from_directory, chunk_documents
    print("Successfully imported modules")
except ImportError as e:
    print(f"Error importing modules: {e}")
    sys.exit(1)

def main():
    # Path to your documents directory
    DOCS_DIR = Path('./docs')
    DOCS_DIR.mkdir(exist_ok=True)
    
    # Load documents
    print(f"Loading documents from {DOCS_DIR}")
    try:
        documents = load_documents_from_directory(DOCS_DIR)
        print(f"Loaded {len(documents)} documents")
    except Exception as e:
        print(f"Error loading documents: {e}")
        return
    
    if not documents:
        print("No documents found. Please add some documents to the 'docs' folder.")
        return
    
    # Display sample document
    if documents:
        print("\nSample document:")
        print(json.dumps(documents[0]['metadata'], indent=2))
    
    # Chunk size and overlap
    CHUNK_SIZE = 1000
    CHUNK_OVERLAP = 200
    
    # Chunk documents
    print(f"\nChunking {len(documents)} documents...")
    try:
        chunked_docs = chunk_documents(
            documents,
            chunk_size=CHUNK_SIZE,
            chunk_overlap=CHUNK_OVERLAP
        )
        print(f"Created {len(chunked_docs)} chunks from {len(documents)} documents")
    except Exception as e:
        print(f"Error chunking documents: {e}")
        return
    
    # Display sample chunk
    if chunked_docs:
        print("\nSample chunk:")
        print(f"Text length: {len(chunked_docs[0]['text'])}")
        print(f"Metadata: {json.dumps(chunked_docs[0]['metadata'], indent=2)}")
    
    # Initialize the retriever
    print("\nInitializing retriever...")
    try:
        retriever = LocalRetriever(persist_dir="./vector_store")
        print("Retriever initialized successfully")
    except Exception as e:
        print(f"Error initializing retriever: {e}")
        return
    
    # Add documents to the retriever
    print("\nAdding documents to the vector store...")
    try:
        retriever.add_documents(chunked_docs)
        print("\nVector store created successfully!")
        print(f"Total documents in the vector store: {len(retriever.documents) if hasattr(retriever, 'documents') else 0}")
    except Exception as e:
        print(f"Error adding documents to vector store: {e}")
        return
    
    # Test the retriever
    print("\nTesting the retriever...")
    try:
        test_queries = [
            "What is this document about?",
            "What are the features of the Mini RAG application?",
            "How does the application work?"
        ]
        
        for query in test_queries:
            print(f"\nQuery: {query}")
            results = retriever.search(query, k=3)
            for i, doc in enumerate(results):
                print(f"\n--- Result {i+1} (Score: {doc.get('score', 0):.4f}) ---")
                print(f"Source: {doc.get('metadata', {}).get('source', 'N/A')}")
                print(f"Text: {doc['text'][:200]}...")
    except Exception as e:
        print(f"Error testing retriever: {e}")

if __name__ == "__main__":
    main()