scripts/ingest_sample.py

#!/usr/bin/env python3
"""
scripts/ingest_sample.py

Ingest sample documents to test the RAG pipeline end-to-end.
Run from the project root: python scripts/ingest_sample.py
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dotenv import load_dotenv
load_dotenv()

from backend.ingestion.web_loader import WebLoader
from backend.ingestion.chunker import SmartChunker
from backend.retrieval.vector_store import get_vector_store
from backend.retrieval.bm25_retriever import get_bm25_retriever

SAMPLE_URLS = [
    "https://en.wikipedia.org/wiki/Retrieval-augmented_generation",
    "https://en.wikipedia.org/wiki/Large_language_model",
    "https://en.wikipedia.org/wiki/Transformer_(deep_learning_architecture)",
]

SAMPLE_TEXTS = [
    {
        "title": "RAG Overview",
        "text": """
Retrieval-Augmented Generation (RAG) is an AI framework that improves the quality of
responses generated by large language models (LLMs) by grounding them in external
knowledge sources. Instead of relying solely on parametric knowledge baked into model
weights, RAG retrieves relevant documents at query time and provides them as context.

The main components of a RAG system are:
1. Document ingestion: Loading, chunking, and embedding documents into a vector store.
2. Retrieval: Using semantic search (and optionally keyword search) to find relevant chunks.
3. Generation: Providing retrieved context to an LLM to generate a grounded answer.

Key advantages of RAG:
- Reduces hallucination by grounding answers in retrieved facts
- Allows knowledge updates without retraining the model
- Enables citation of sources for transparency
- More cost-effective than fine-tuning for knowledge injection
        """,
    },
    {
        "title": "Hybrid Search Explained",
        "text": """
Hybrid search combines dense (semantic) retrieval with sparse (keyword) retrieval.
Dense retrieval uses embedding models to encode text into high-dimensional vectors,
enabling semantic similarity matching. Sparse retrieval uses algorithms like BM25
to match based on term frequency and inverse document frequency.

Reciprocal Rank Fusion (RRF) is the standard method for combining ranked lists from
multiple retrieval systems. The formula is: RRF(d) = sum(1 / (k + rank(d))) where
k=60 is a smoothing constant. RRF is rank-based (not score-based), making it robust
to score distribution differences between retrieval systems.

Re-ranking with cross-encoders (like Cohere Rerank) further improves precision by
scoring each query-document pair jointly, rather than independently. This is the
retrieve-then-rerank pattern used by production systems at Notion, Perplexity, and You.com.
        """,
    },
]


def ingest_urls():
    print("\n=== Ingesting URLs ===")
    loader = WebLoader()
    vs = get_vector_store()
    bm25 = get_bm25_retriever()

    for url in SAMPLE_URLS:
        print(f"  Fetching: {url}")
        try:
            chunks = loader.load_url_sync(url)
            vs.add_chunks(chunks)
            bm25.add_chunks(chunks)
            print(f"  ✓ Ingested {len(chunks)} chunks")
        except Exception as e:
            print(f"  ✗ Failed: {e}")


def ingest_texts():
    print("\n=== Ingesting Sample Texts ===")
    chunker = SmartChunker()
    vs = get_vector_store()
    bm25 = get_bm25_retriever()

    import hashlib
    for item in SAMPLE_TEXTS:
        source_id = "txt_" + hashlib.sha256(item["title"].encode()).hexdigest()[:12]
        chunks = chunker.chunk_text(
            text=item["text"],
            source_id=source_id,
            source_type="text",
            source_name=item["title"],
        )
        vs.add_chunks(chunks)
        bm25.add_chunks(chunks)
        print(f"  ✓ '{item['title']}' → {len(chunks)} chunks")


def test_query():
    print("\n=== Testing Query Pipeline ===")
    from backend.agents.research_agent import get_research_agent
    agent = get_research_agent()

    query = "What is hybrid search and why is it better than pure semantic search?"
    print(f"  Query: {query}")

    result = agent.query(query)
    print(f"\n  Answer:\n{result.answer[:500]}...")
    print(f"\n  Citations:")
    for c in result.citations[:3]:
        print(f"    - {c.source_name} (score: {c.relevance_score:.3f})")


if __name__ == "__main__":
    print("RAG Research Agent — Sample Ingestion Script")
    print("=" * 50)

    ingest_texts()   # Always works (no network needed)
    ingest_urls()    # Requires internet

    test_query()     # Requires OPENAI_API_KEY

    print("\n✅ Done! Your knowledge base is ready.")
    print("   Start the API: uvicorn backend.api.main:app --reload")