ingest.py

"""
PaceWise Data Ingestion Script
Processes running_data.json and creates ChromaDB vector store
"""

import json
import os
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma 

def load_running_data(json_file='running_data.json'):
    """Load running discussion data from JSON file"""
    print(f"Loading data from {json_file}...")
    with open(json_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    print(f"Loaded {len(data)} threads")
    return data

def prepare_documents(threads):
    """
    Convert threads into document format for vector store
    Each document contains: thread context + individual comment
    """
    documents = []
    metadatas = []
    
    for thread in threads:
        thread_id = thread['thread_id']
        subreddit = thread['subreddit']
        title = thread['title']
        post_text = thread['post_text']
        
        # Create context header for each thread
        thread_context = f"Thread: {title}\nSubreddit: r/{subreddit}\nQuestion: {post_text}\n\n"
        
        # Process each comment as a separate document with thread context
        for idx, comment in enumerate(thread['comments']):
            comment_text = comment['text']
            
            # Combine thread context with comment
            full_text = thread_context + f"Response: {comment_text}"
            
            documents.append(full_text)
            
            # Store metadata for citation and filtering
            metadatas.append({
                'thread_id': thread_id,
                'subreddit': subreddit,
                'title': title,
                'post_text': post_text,
                'comment_index': idx,
                'score': comment.get('score', 0)
            })
    
    print(f"Prepared {len(documents)} documents from comments")
    return documents, metadatas

def create_vector_store(documents, metadatas, persist_directory='chroma_db'):
    """
    Create ChromaDB vector store with embeddings
    Uses HuggingFace sentence-transformers for embeddings
    """
    print("Initializing embeddings model...")
    
    # Use HuggingFace sentence-transformers (same as PulsePoint)
    # This model is lightweight and works well for semantic search
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2",
        model_kwargs={'device': 'cpu'},  # Use CPU for compatibility
        encode_kwargs={'normalize_embeddings': True}
    )
    
    print("Creating text chunks for embedding...")
    
    # Don't chunk - each comment with thread context is already appropriately sized
    split_docs = documents
    split_metadatas = metadatas
    
    print(f"Created {len(split_docs)} text chunks")
    
    # Create ChromaDB vector store
    print("Creating vector store and generating embeddings...")
    print("(This may take a few minutes...)")
    
    vectorstore = Chroma.from_texts(
        texts=split_docs,
        embedding=embeddings,
        metadatas=split_metadatas,
        persist_directory=persist_directory
    )
    
    # Persist to disk
    vectorstore.persist()
    
    print(f"Vector store created and saved to {persist_directory}/")
    print(f"Total embedded chunks: {len(split_docs)}")
    
    return vectorstore

def verify_vector_store(persist_directory='chroma_db'):
    """Verify the vector store was created successfully"""
    print("\nVerifying vector store...")
    
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2",
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True}
    )
    
    # Load the persisted vector store
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embeddings
    )
    
    # Test query with MMR for diversity
    test_query = "How should I build mileage for marathon training?"
    results = vectorstore.max_marginal_relevance_search(test_query, k=3, fetch_k=10)
    
    print(f"\nTest query: '{test_query}'")
    print(f"Retrieved {len(results)} results")
    
    for i, doc in enumerate(results, 1):
        print(f"\n--- Result {i} ---")
        print(f"Thread: {doc.metadata.get('title', 'N/A')}")
        print(f"Subreddit: r/{doc.metadata.get('subreddit', 'N/A')}")
        print(f"Preview: {doc.page_content[:200]}...")
    
    print("\n✓ Vector store verification complete!")

def main():
    """Main ingestion pipeline"""
    print("="*60)
    print("PaceWise Data Ingestion Pipeline")
    print("="*60)
    
    # Step 1: Load data
    threads = load_running_data('running_data.json')
    
    # Step 2: Prepare documents
    documents, metadatas = prepare_documents(threads)
    
    # Step 3: Create vector store
    persist_dir = 'chroma_db'
    
    # Remove existing vector store if it exists
    if os.path.exists(persist_dir):
        print(f"\nRemoving existing vector store at {persist_dir}/")
        import shutil
        shutil.rmtree(persist_dir)
    
    vectorstore = create_vector_store(documents, metadatas, persist_dir)
    
    # Step 4: Verify
    verify_vector_store(persist_dir)
    
    print("\n" + "="*60)
    print("Ingestion complete! Vector store ready for RAG queries.")
    print("="*60)

if __name__ == "__main__":
    main()