example_rag.py

# ruff: noqa: F401
import os
import sys
from datetime import datetime as dt
from pathlib import Path

import click
import typer
from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain.messages import AIMessage, HumanMessage, SystemMessage
from langchain_classic.storage import LocalFileStore
from langchain_community.docstore.in_memory import InMemoryDocstore
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter
from rich import print as rprint
from rich.console import Console
from rich.markdown import Markdown
from rich.pretty import pprint

# import faiss

# --- CONFIG
load_dotenv()

DOCS_DIR = Path("rag-docs")
# DOCS_DIR = Path("python-3.14-docs-text")
VECTOR_DB_DIR = Path("vector_store")
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
EMBED_MODEL_KWARGS = {"device": "cpu"}
# EMBED_MODEL_KWARGS = {"device": "cuda"}

# ---

app = typer.Typer(pretty_exceptions_show_locals=False, pretty_exceptions_short=True)


def get_llm_openai():
    return ChatOpenAI(
        model="gpt-4o-mini",
        temperature=0,
    )


def get_llm_groq(model):
    return ChatGroq(model=model)


def get_llm_openrouter(model):
    llm = ChatOpenAI(
        model=model,
        openai_api_key=os.getenv("OPENROUTER_KEY"),
        openai_api_base="https://openrouter.ai/api/v1",
        # default_headers={
        #     "HTTP-Referer": "https://your-site-url.com", # Optional: for OpenRouter rankings
        #     "X-Title": "Your App Name",                 # Optional: for OpenRouter rankings
        # }
    )
    return llm


def get_llm_gemini(model):
    return ChatGoogleGenerativeAI(
        model=model,
        temperature=1.0,  # Gemini 3.0+ defaults to 1.0
        max_tokens=None,
        timeout=None,
    )


def load_documents():
    docs = []
    for path in DOCS_DIR.glob("**/*.txt"):
        docs.append(
            # metadata={"source": str(path)})
            Document(
                page_content=path.read_text(encoding="utf8"),
                metadata={"source": path.relative_to(DOCS_DIR).as_posix()},
            )
        )
    if len(docs) == 0:
        raise FileNotFoundError(f"Empty directory: {DOCS_DIR.as_posix()}")

    return docs


def get_vector_store(docs):
    embeddings = HuggingFaceEmbeddings(
        model_name=EMBED_MODEL, model_kwargs=EMBED_MODEL_KWARGS
    )

    if VECTOR_DB_DIR.exists():
        print("Loading FAISS vector store from file...")
        return FAISS.load_local(
            VECTOR_DB_DIR, embeddings, allow_dangerous_deserialization=True
        )

    if docs is None:
        raise ValueError("get_vector_store: docs is None")
    elif len(docs) == 0:
        raise ValueError(f"get_vector_store: {len(docs)=}")

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100,
    )
    chunks = splitter.split_documents(docs)

    print("Creating vector store...")
    tic = dt.now()
    db = FAISS.from_documents(chunks, embeddings)
    tac = dt.now()
    print(f"Took {tac - tic}")

    db.save_local(VECTOR_DB_DIR)
    print(f"Saved vector store to {VECTOR_DB_DIR}")
    return db


@app.command()
def embed():
    """
    Embeds all documents and creates a vector store.

    I recommend setting EMBED_MODEL_KWARGS = {"device": "cuda"}
    """
    if not DOCS_DIR.exists():
        print(f"Create {DOCS_DIR} directory with .txt files to start.")
        sys.exit(1)

    docs = load_documents()
    rprint(f"Found {len(docs)} documents in {DOCS_DIR.as_posix()}")
    print("Using embedding model ", EMBED_MODEL)
    print(" - kwargs: ", EMBED_MODEL_KWARGS)
    db = get_vector_store(docs)

    print(db)


# ---

PROMPT = ChatPromptTemplate.from_template("""
You are a helpful assistant.
Answer the question using ONLY the context below.
If the answer is not in the context, say you don't know.

Context:
{context}

Question:
{question}
""")


def ask_rag(db, llm, query, *, k=4, verbose=False):
    docs = db.similarity_search(query, k=k)
    context = "\n\n".join(d.page_content for d in docs)

    if verbose:
        print("--- Fetched chunks metadata:")
        # currently: {'source': 'whatsnew/3.14.txt'}
        pprint(
            [d.metadata for d in docs],
            indent_guides=False,
        )
        print("--- context:")
        pprint(
            [d.page_content for d in docs],
            max_string=80,
            max_depth=2,
            indent_guides=False,
        )
        print("---\n")

    chain = PROMPT | llm
    return docs, chain.invoke(
        {
            "context": context,
            "question": query,
        }
    ).content


@app.command()
def chat(verbose: bool = False, show_sources_url: bool = True):
    """
    Ask questions about the documents.

    Assumes the vector store is already prepared.

    Set --verbose to see retrieved chunks.
    """
    console = Console()
    if not DOCS_DIR.exists():
        print(f"Create `{DOCS_DIR}` directory with .txt files to start.")
        sys.exit(1)

    docs = load_documents()
    print(f"Found {len(docs)} documents")
    rprint("Using embedding model ", EMBED_MODEL)
    db = get_vector_store(docs)

    # llm = get_llm_gemini("gemini-flash-latest")
    # llm = get_llm_gemini("gemini-flash-lite-latest")
    llm = get_llm_groq("openai/gpt-oss-120b")
    # llm = get_llm_openrouter("nvidia/nemotron-3-nano-30b-a3b:free")
    if hasattr(llm, "model_name"):
        rprint(f"Chatting with LLM {llm.model_name}")
    elif hasattr(llm, "model"):
        rprint(f"Chatting with LLM {llm.model}")
    else:
        pprint(llm)

    print("RAG CLI ready. Ctrl+C to exit.\n")

    while True:
        try:
            # query = input(" Question > ")
            query = click.prompt("Question")
            query = query.strip()
            if len(query) < 2:
                break
        except KeyboardInterrupt:
            break

        rag_docs, answer = ask_rag(db, llm, query, k=10, verbose=verbose)
        print("")
        try:
            console.print(Markdown(answer))
        except Exception as e:
            print("Cannot show markdown, ", type(e))
            rprint(answer)
        print("")

        if show_sources_url:
            print("Sources:")
            rag_sources_url = [each_doc.metadata.get("source") for each_doc in rag_docs]
            rag_sources_url = list(set(rag_sources_url))
            rag_sources_url.sort()
            for relpath in rag_sources_url:
                relpath = relpath.replace(".txt", ".html")
                print(f" - https://docs.python.org/3/{relpath}")

    print("\n\nBye!")


if __name__ == "__main__":
    # import mlflow
    # print("mlflow autologging")
    # mlflow.langchain.autolog()
    app()