RAG Tutorial

# Step 1: Install dependencies
# pip install langchain chromadb openai langchain-community langchain-openai

from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain_openai import ChatOpenAI

# ── INDEXING ──────────────────────────
loader = PyPDFLoader("test_documentation.pdf")
documents = loader.load()

splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200
)
chunks = splitter.split_documents(documents)
print(f"Created {len(chunks)} chunks")

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    persist_directory="./chroma_db"
)
retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

# ── QUERY & GENERATE ──────────────────
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
qa_chain = RetrievalQA.from_chain_type(
    llm=llm, chain_type="stuff",
    retriever=retriever,
    return_source_documents=True
)

response = qa_chain.invoke({"query": "What are the login test cases?"})
print(response["result"])
print("Sources:", [doc.metadata for doc in response["source_documents"]])

from langchain.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_experimental.text_splitter import SemanticChunker
from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank

# ── HyDE: Query → Hypothetical Doc → Better Retrieval ──
hyde_prompt = ChatPromptTemplate.from_template(
    "Write a short passage that would answer: {question}"
)
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
hyde_chain = hyde_prompt | llm

hypothetical_doc = hyde_chain.invoke({"question": "How to test API rate limiting?"})
# Now use hypothetical_doc.content for embedding + retrieval

# ── Semantic Chunking ─────────────────────────────────
embeddings = OpenAIEmbeddings()
semantic_splitter = SemanticChunker(
    embeddings,
    breakpoint_threshold_type="percentile",
    breakpoint_threshold_amount=95
)
semantic_chunks = semantic_splitter.split_documents(documents)

# ── Cohere Re-ranking: top-20 → best 4 ───────────────
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})
reranker = CohereRerank(model="rerank-english-v3.0", top_n=4)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=base_retriever
)

results = compression_retriever.invoke(
    "Edge cases for payment testing?"
)
print(f"Re-ranked: {len(results)} top documents returned")

from langchain.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain.schema.runnable import RunnableLambda, RunnablePassthrough

# Specialized retrievers per domain
api_retriever = vectorstore_api.as_retriever(search_kwargs={"k": 4})
ui_retriever  = vectorstore_ui.as_retriever(search_kwargs={"k": 4})
perf_retriever= vectorstore_perf.as_retriever(search_kwargs={"k": 4})

# Router LLM (temperature=0 for deterministic classification)
router_prompt = ChatPromptTemplate.from_template(
    """Classify this QA query into ONE of:
    api_testing | ui_testing | performance_testing
    Query: {question}
    Category:"""
)
router_chain = router_prompt | ChatOpenAI(temperature=0)

def route_query(info):
    route = info["route"].content.strip().lower()
    if "api" in route:
        return api_retriever.invoke(info["question"])
    elif "ui" in route:
        return ui_retriever.invoke(info["question"])
    else:
        return perf_retriever.invoke(info["question"])

# Swappable generator pattern
class ModularRAG:
    def __init__(self, retriever, reranker=None, generator=None):
        self.retriever = retriever
        self.reranker  = reranker
        self.generator = generator or ChatOpenAI(model="gpt-4o-mini")

    def query(self, question: str):
        docs = self.retriever.invoke(question)
        if self.reranker:
            docs = self.reranker.compress_documents(docs, question)
        context = "\n".join([d.page_content for d in docs])
        return self.generator.invoke(
            f"Context: {context}\nQuestion: {question}\nAnswer:"
        )

# Easily swap any component!
rag_v1 = ModularRAG(retriever=bm25_retriever, generator=gpt4)
rag_v2 = ModularRAG(retriever=vector_retriever, reranker=cohere, generator=claude)

from langchain_community.graphs import Neo4jGraph
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain.chains import GraphCypherQAChain
from langchain_openai import ChatOpenAI

# Connect to Neo4j graph database
graph = Neo4jGraph(
    url="bolt://localhost:7687",
    username="neo4j",
    password="your_password"
)

# Extract entities + relationships from documents
llm = ChatOpenAI(model="gpt-4o", temperature=0)
transformer = LLMGraphTransformer(llm=llm)

graph_documents = transformer.convert_to_graph_documents(documents)
graph.add_graph_documents(graph_documents)
print(f"Graph built with {len(graph_documents)} document graphs")

# Query the graph with natural language
chain = GraphCypherQAChain.from_llm(
    llm=llm, graph=graph, verbose=True
)

result = chain.invoke({
    "query": "Which test suites cover the payment module?"
})
# LLM generates: MATCH (t:TestSuite)-[:COVERS]->(m:Module {name:'payment'}) RETURN t

# Microsoft GraphRAG for community-based search
# graphrag query --method global --query "Main testing challenges?"
# graphrag query --method local  --query "Auth service test needs?"

from langchain.agents import AgentExecutor, create_react_agent
from langchain_openai import ChatOpenAI
from langchain.tools.retriever import create_retriever_tool
from langchain import hub

# Define specialized retriever tools
api_tool = create_retriever_tool(
    api_retriever, name="api_test_docs",
    description="Search API testing documentation and test cases"
)
ui_tool = create_retriever_tool(
    ui_retriever, name="ui_test_docs",
    description="Search UI testing guides, Selenium and Playwright docs"
)
bug_tool = create_retriever_tool(
    bug_retriever, name="bug_reports",
    description="Search historical bug reports and known issues"
)
tools = [api_tool, ui_tool, bug_tool]

# Create ReAct agent (Reason + Act loop)
llm = ChatOpenAI(model="gpt-4o", temperature=0)
prompt = hub.pull("hwchase17/react")
agent = create_react_agent(llm, tools, prompt)

agent_executor = AgentExecutor(
    agent=agent, tools=tools, verbose=True,
    max_iterations=5,           # Prevent infinite loops!
    handle_parsing_errors=True
)

# Multi-tool query: agent decides which tools to use
response = agent_executor.invoke({
    "input": "Were there payment API bugs last quarter, "
             "and what test cases should cover them?"
})
# Agent reasoning trace:
# Thought: I need both bug history and API test cases
# Action: bug_reports("payment API Q3")
# Observation: Found 3 bugs: AUTH-401, TIMEOUT-503...
# Action: api_test_docs("payment endpoint testing")
# Final Answer: Based on bugs found, recommend...

from langgraph.graph import StateGraph, END
from typing import TypedDict, List

class RAGState(TypedDict):
    question: str
    documents: List[str]
    generation: str
    needs_retrieval: bool
    is_relevant: bool
    is_grounded: bool

def decide_retrieval(state):
    # Reflection token: [Retrieve]
    prompt = f'Does "{state["question"]}" need external info? (yes/no)'
    response = llm.invoke(prompt)
    return {"needs_retrieval": "yes" in response.content.lower()}

def check_relevance(state):
    # Reflection token: [ISREL]
    prompt = f'Are these docs relevant to "{state["question"]}"?\n'
    prompt += "\n".join(state["documents"][:2])
    response = llm.invoke(prompt)
    return {"is_relevant": "yes" in response.content.lower()}

def check_groundedness(state):
    # Reflection token: [ISSUP]
    prompt = f"Is this answer grounded in the context?\n"
    prompt += f'Answer: {state["generation"]}\nContext: {state["documents"][0]}'
    response = llm.invoke(prompt)
    return {"is_grounded": "yes" in response.content.lower()}

# Build the LangGraph state machine
workflow = StateGraph(RAGState)
workflow.add_node("decide",    decide_retrieval)
workflow.add_node("retrieve",  retrieve)
workflow.add_node("relevance", check_relevance)
workflow.add_node("generate",  generate)
workflow.add_node("grounded",  check_groundedness)

workflow.set_entry_point("decide")
workflow.add_conditional_edges(
    "decide", lambda s: "retrieve" if s["needs_retrieval"] else "generate"
)
workflow.add_edge("retrieve", "relevance")
workflow.add_conditional_edges(
    "relevance", lambda s: "generate" if s["is_relevant"] else "retrieve"
)
workflow.add_edge("generate", "grounded")
workflow.add_conditional_edges(
    "grounded", lambda s: END if s["is_grounded"] else "retrieve"
)
app = workflow.compile()

from langgraph.graph import StateGraph, END
from langchain_community.tools.tavily_search import TavilySearchResults

web_search = TavilySearchResults(max_results=3)

def grade_documents(state):
    "Grade each retrieved document 0.0-1.0 for relevance"
    question = state["question"]
    scored_docs = []

    for doc in state["documents"]:
        grade_prompt = f"""Grade relevance 0.0-1.0:
        Question: {question}
        Document: {doc.page_content[:500]}
        Score:"""
        score = float(llm.invoke(grade_prompt).content.strip())
        scored_docs.append((doc, score))

    avg_score = sum(s for _, s in scored_docs) / len(scored_docs)

    if   avg_score > 0.7: action = "correct"
    elif avg_score > 0.4: action = "ambiguous"
    else:                action = "incorrect"

    relevant = [d for d, s in scored_docs if s > 0.5]
    return {"action": action, "documents": relevant}

def web_search_fallback(state):
    "Fallback to web search when retrieval confidence is low"
    results = web_search.invoke({"query": state["question"]})
    web_docs = [r["content"] for r in results]
    existing = state.get("documents", [])
    return {"documents": existing + web_docs}

def route_after_grading(state):
    if   state["action"] == "correct": return "generate"
    else:                               return "web_search"  # ambiguous or incorrect

workflow = StateGraph(CRAGState)
workflow.add_edge("retrieve", "grade")
workflow.add_conditional_edges("grade", route_after_grading)
workflow.add_edge("web_search", "generate")
workflow.add_edge("generate", END)
crag_app = workflow.compile()

from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever

# BM25: keyword search on raw text
bm25_retriever = BM25Retriever.from_documents(
    documents=chunks, k=4
)

# Vector: semantic search on embeddings
vector_retriever = vectorstore.as_retriever(
    search_kwargs={"k": 4}
)

# Ensemble: weights favor semantic (0.6) over keyword (0.4)
hybrid_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.4, 0.6]
)

results = hybrid_retriever.invoke("Error ERR_AUTH_403 during login test")
# BM25 finds exact error code matches
# Vector finds semantically related auth failure docs
# Results merged and deduplicated automatically

# ── Reciprocal Rank Fusion (manual) ──────────────────
def reciprocal_rank_fusion(results_list, k=60):
    fused_scores = {}
    for results in results_list:
        for rank, doc in enumerate(results):
            doc_id = doc.page_content[:100]
            if doc_id not in fused_scores:
                fused_scores[doc_id] = {"doc": doc, "score": 0}
            fused_scores[doc_id]["score"] += 1.0 / (rank + k)

    sorted_docs = sorted(
        fused_scores.values(),
        key=lambda x: x["score"], reverse=True
    )
    return [item["doc"] for item in sorted_docs]

bm25_r  = bm25_retriever.invoke(query)
vector_r = vector_retriever.invoke(query)
hybrid_r = reciprocal_rank_fusion([bm25_r, vector_r])

from langchain.retrievers.multi_vector import MultiVectorRetriever
from langchain.storage import InMemoryByteStore
from langchain_openai import ChatOpenAI
import uuid

vision_llm = ChatOpenAI(model="gpt-4o", temperature=0)

def summarize_image(image_base64: str) -> str:
    "Use Vision LLM to describe image for a QA engineer"
    response = vision_llm.invoke([{
        "type": "image_url",
        "image_url": {"url": f"data:image/png;base64,{image_base64}"}
    }, {
        "type": "text",
        "text": "Describe for a QA engineer: error messages, "
               "UI elements, test results, and status codes visible."
    }])
    return response.content

# Multi-Vector store: summaries → vector store, originals → byte store
byte_store = InMemoryByteStore()
id_key = "doc_id"

multi_retriever = MultiVectorRetriever(
    vectorstore=vectorstore,
    byte_store=byte_store,
    id_key=id_key
)

# Index each image
for img_path in image_paths:
    img_b64 = encode_image(img_path)
    summary = summarize_image(img_b64)
    doc_id = str(uuid.uuid4())

    # Summary → vector store (for retrieval)
    multi_retriever.vectorstore.add_documents([
        Document(page_content=summary, metadata={id_key: doc_id})
    ])
    # Original image → byte store (for return)
    multi_retriever.docstore.mset([(doc_id, img_b64)])

# Query retrieves summaries but RETURNS original images
results = multi_retriever.invoke("Login page 404 error screenshot")

from langchain_openai import ChatOpenAI
from langchain.schema import Document

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

CONTEXT_PROMPT = """Here is the full document:
<document>
{document}
</document>

Here is a chunk from that document:
<chunk>
{chunk}
</chunk>

Provide 2-3 sentence context explaining:
1. What document/section this chunk belongs to
2. What the chunk is specifically about
3. Key entities or identifiers mentioned

Context:"""

def add_context_to_chunks(full_doc: str, chunks: list) -> list:
    "Prepend document context to each chunk before embedding"
    contextualized = []
    for chunk in chunks:
        context = llm.invoke(
            CONTEXT_PROMPT.format(
                document=full_doc[:4000],
                chunk=chunk.page_content
            )
        ).content

        enriched = f"{context}\n\n{chunk.page_content}"
        contextualized.append(Document(
            page_content=enriched,
            metadata={**chunk.metadata, "original": chunk.page_content}
        ))
    return contextualized

# ── Best combination: Contextual + Hybrid + Rerank ───
ctx_chunks = add_context_to_chunks(full_doc_text, chunks)

ctx_bm25 = BM25Retriever.from_documents(ctx_chunks, k=20)
ctx_vector = Chroma.from_documents(ctx_chunks, embeddings).as_retriever(
    search_kwargs={"k": 20}
)
hybrid = EnsembleRetriever(
    retrievers=[ctx_bm25, ctx_vector], weights=[0.4, 0.6]
)
final = ContextualCompressionRetriever(
    base_compressor=CohereRerank(top_n=5),
    base_retriever=hybrid
)
# This combination: -67% retrieval failures (Anthropic research)