query.py

import os
from dotenv import load_dotenv

from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

load_dotenv()

GROQ_API_KEY    = os.getenv("GROQ_API_KEY")
GROQ_MODEL      = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "all-MiniLM-L6-v2")
TOP_K           = int(os.getenv("RETRIEVAL_TOP_K", "4"))
CHROMA_DIR      = "./chroma_db"
COLLECTION_NAME = "my_documents"


# RAG PROMPT
RAG_PROMPT = """You are a helpful assistant that answers questions based ONLY
on the document excerpts provided below.

────────────────────────────────
CONTEXT FROM YOUR DOCUMENTS:
{context}
────────────────────────────────

QUESTION: {question}

ANSWER:"""


# Helper: format retrieved chunks for the prompt
def format_retrieved_chunks(docs: list) -> str:
    parts = []
    for i, doc in enumerate(docs, 1):
        source = os.path.basename(doc.metadata.get("source", "unknown"))
        page   = doc.metadata.get("page", "?")
        parts.append(
            f"[Excerpt {i} | File: {source} | Page: {page}]\n{doc.page_content}"
        )
    return "\n\n---\n\n".join(parts)


# Build the RAG chain
def build_chain(vector_store: Chroma):
    """
    Create the full RAG pipeline using LangChain Expression Language (LCEL).

    """
    retriever = vector_store.as_retriever(
        search_type="similarity",
        search_kwargs={"k": TOP_K},
    )

    prompt = ChatPromptTemplate.from_template(RAG_PROMPT)

    # Groq LLM 
    llm = ChatGroq(
        model=GROQ_MODEL,
        groq_api_key=GROQ_API_KEY,
        temperature=0,        # 0 = deterministic answers, less hallucination
        max_tokens=1024,
    )

    # The full chain:
    #   { context: retrieve→format, question: passthrough }
    #      → prompt
    #      → LLM
    #      → parse to string
    chain = (
        {
            "context":  retriever | format_retrieved_chunks,
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm
        | StrOutputParser()
    )

    return chain, retriever


# Interactive Q&A loop
def main():
    print("\n" + "=" * 55)
    print("  RAG — Ask Questions")
    print("=" * 55)

    # Load embedding model
    print("\nLoading embedding model...")
    embedding_fn = HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )

    # Connect to existing ChromaDB
    print("Connecting to ChromaDB...")
    vector_store = Chroma(
        collection_name=COLLECTION_NAME,
        embedding_function=embedding_fn,
        persist_directory=CHROMA_DIR,
    )

    count = vector_store._collection.count()
    if count == 0:
        print("\n  ERROR: ChromaDB is empty.")
        print("  Run 'uv run python index.py' first.\n")
        return

    print(f"  Ready! {count} chunks loaded from ChromaDB.")
    print(f"  Using model: {GROQ_MODEL}")
    print("\n  Commands: 'quit' to exit | 'sources' to toggle source display")

    # Build the RAG chain
    chain, retriever = build_chain(vector_store)
    show_sources = True

    # Q&A loop
    while True:
        print("\n" + "─" * 55)
        try:
            question = input("  Your question: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\n  Goodbye!")
            break

        if not question:
            continue
        if question.lower() in ("quit", "exit", "q"):
            print("  Goodbye!")
            break
        if question.lower() == "sources":
            show_sources = not show_sources
            status = "ON" if show_sources else "OFF"
            print(f"  Source display toggled {status}")
            continue

        # --- Retrieve and answer ---
        print("\n  Searching documents and generating answer...\n")

        # Get the retrieved chunks (for showing sources)
        retrieved = retriever.invoke(question)

        # Run the full RAG chain
        answer = chain.invoke(question)

        # Display answer
        print(f"  ANSWER:\n")
        # Print answer with word wrap at ~70 chars
        words = answer.split()
        line = "  "
        for word in words:
            if len(line) + len(word) > 72:
                print(line)
                line = "  " + word + " "
            else:
                line += word + " "
        if line.strip():
            print(line)

        # Display sources
        if show_sources and retrieved:
            print(f"\n  SOURCES ({len(retrieved)} chunks retrieved):")
            for i, doc in enumerate(retrieved, 1):
                src  = os.path.basename(doc.metadata.get("source", "unknown"))
                page = doc.metadata.get("page", "?")
                snippet = doc.page_content[:100].replace("\n", " ")
                print(f"\n  [{i}] {src} — Page {page}")
                print(f"      \"{snippet}...\"")


if __name__ == "__main__":
    main()