main.py

"""
ChatJio — end-to-end RAG pipeline entry point.

Flow:
  User query
      → retrieve()        [retrieval/retriever.py]   embed + hybrid search + rerank
      → Generator.generate()  [generation/generator.py]  prompt + LLM
      → answer

Usage:
  python main.py                         # interactive chat loop
  python main.py --query "your question" # single query
  python main.py --ingest --url https://www.jioinstitute.edu.in/
  python main.py --ingest --data-dir data/raw
"""

import argparse
import logging
import sys

logging.basicConfig(level=logging.WARNING, format="%(levelname)s | %(name)s | %(message)s")

from retrieval.retriever import retrieve
from generation.generator import Generator, GenerationError
from ingestion.pipeline import run_ingestion
from retrieval.vector_store import upsert_chunks


def answer(query: str, generator: Generator) -> dict:
    """Single-shot answer used by --query mode. No interactive confirmation."""
    chunks, match_status = retrieve(query)

    if match_status == "no_match":
        print("(No match found in DB — falling back to LLM general knowledge.)")
        return generator.generate(query, [])

    if match_status == "partial":
        return generator.generate(query, [])

    return generator.generate(query, chunks)


def run_chat(generator: Generator):
    print("ChatJio — Ask anything about Jio Institute. Type 'exit' to quit.\n")
    while True:
        try:
            query = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if not query:
            continue
        if query.lower() in ("exit", "quit", "q"):
            print("Goodbye!")
            break

        try:
            chunks, match_status = retrieve(query)

            if match_status == "good":
                result = generator.generate(query, chunks)
                print(f"\nChatJio: Found it in the DB.\n")

            else:
                label = "no results found" if match_status == "no_match" else "partial results found"
                print(
                    f"\nChatJio: Sorry, {label} on this topic in the stored DB. "
                    "Do you want me to fetch it outside the DB? (yes/no)"
                )
                try:
                    confirm = input("You: ").strip().lower()
                except (EOFError, KeyboardInterrupt):
                    print("\nGoodbye!")
                    break
                if confirm not in ("yes", "y"):
                    print("\nChatJio: Okay, skipping this one.\n")
                    continue
                result = generator.generate(query, [])
                print(f"\nChatJio: No results found in DB, seeking external help.\n")

            print(f"\nChatJio: {result['answer']}")
            if result["sources"]:
                unique_sources = list(dict.fromkeys(result["sources"]))
                print(f"Sources: {', '.join(unique_sources)}")
            print()

        except GenerationError as e:
            print(f"\n[Error] {e}\n")


def run_ingest(data_dir: str = None, url: str = None, max_pages: int = 1000):
    print("Starting ingestion...")
    chunks = run_ingestion(data_dir=data_dir, url=url, max_pages=max_pages)
    print(f"Ingestion complete — {len(chunks)} chunks produced.")
    print("Embedding and storing...")
    from ingestion.embedder import embed_chunks
    chunks = embed_chunks(chunks)
    stored = upsert_chunks(chunks)
    print(f"Done. {stored} vectors stored in Qdrant.")


def main():
    parser = argparse.ArgumentParser(description="ChatJio RAG pipeline")
    parser.add_argument("--query", type=str, help="Run a single query and exit")
    parser.add_argument("--ingest", action="store_true", help="Run ingestion instead of chat")
    parser.add_argument("--url", type=str, help="Website URL to ingest")
    parser.add_argument("--data-dir", type=str, help="Local directory of PDFs to ingest")
    parser.add_argument("--max-pages", type=int, default=1000, help="Max pages to crawl (default 1000)")
    args = parser.parse_args()

    if args.ingest:
        if not args.url and not args.data_dir:
            print("Error: --ingest requires --url or --data-dir")
            sys.exit(1)
        run_ingest(data_dir=args.data_dir, url=args.url, max_pages=args.max_pages)
        return

    generator = Generator()

    if args.query:
        result = answer(args.query, generator)
        print(f"\nAnswer: {result['answer']}")
        if result["sources"]:
            unique_sources = list(dict.fromkeys(result["sources"]))
            print(f"Sources: {', '.join(unique_sources)}")
    else:
        run_chat(generator)


if __name__ == "__main__":
    main()