Document RAG

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

# Stores: vector index for search, byte store for original docs
vectorstore = FAISS.from_documents([], OpenAIEmbeddings())
docstore = InMemoryByteStore()
id_key = "doc_id"

retriever = MultiVectorRetriever(
    vectorstore=vectorstore,
    byte_store=docstore,
    id_key=id_key,
)

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

for chunk in chunks:
    doc_id = str(uuid.uuid4())
    chunk.metadata[id_key] = doc_id

    # Store original chunk in docstore
    docstore.mset([(doc_id, chunk)])

    # Representation 1: raw text embedding
    vectorstore.add_documents([chunk])

    # Representation 2: LLM summary
    summary = llm.invoke(f"Summarize in 2 sentences:\n{chunk.page_content}")
    summary_doc = Document(
        page_content=summary.content,
        metadata={id_key: doc_id, "type": "summary"}
    )
    vectorstore.add_documents([summary_doc])

    # Representation 3: hypothetical questions (HyDE variant)
    questions = llm.invoke(
        f"Generate 3 questions this text answers:\n{chunk.page_content}"
    )
    for q in questions.content.split("\n"):
        if q.strip():
            q_doc = Document(
                page_content=q.strip(),
                metadata={id_key: doc_id, "type": "question"}
            )
            vectorstore.add_documents([q_doc])

# At query time: user question matches a hypothetical question,
# but the LLM receives the original full chunk
results = retriever.invoke("What was the revenue in Q3?")
# Returns the original financial table, not the summary or question

from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

def generate_with_citations(query: str, chunks: list) -> dict:
    """Generate answer with inline citations and verification."""

    # Format chunks with citation markers and rich metadata
    context_parts = []
    for i, chunk in enumerate(chunks):
        source = chunk.metadata.get("source", "Unknown")
        page = chunk.metadata.get("page", "N/A")
        section = chunk.metadata.get("h2", chunk.metadata.get("section", ""))
        chunk_type = chunk.metadata.get("chunk_type", "text")

        header = f"[{i+1}] Source: {source} | Page: {page}"
        if section:
            header += f" | Section: {section}"
        if chunk_type == "table":
            header += " | [TABLE]"

        context_parts.append(f"{header}\n{chunk.page_content}")

    context = "\n\n---\n\n".join(context_parts)

    prompt = ChatPromptTemplate.from_messages([
        ("system", """You are a precise document analyst. Rules:
1. Answer ONLY using the provided context chunks.
2. Include inline citations [1], [2], etc. for EVERY claim.
3. If a claim spans multiple sources, cite all: [1][3].
4. For tables, reference the specific row/column.
5. If the context is insufficient, say "The provided documents
   do not contain enough information to answer this question."
6. Never infer beyond what the sources explicitly state."""),
        ("user", """Context chunks:
{context}

Question: {question}

Answer with citations:""")
    ])

    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    chain = prompt | llm

    response = chain.invoke({
        "context": context,
        "question": query
    })

    # Extract and validate citations
    import re
    cited_nums = set(int(n) for n in re.findall(r'\[(\d+)\]', response.content))
    available_nums = set(range(1, len(chunks) + 1))
    hallucinated_citations = cited_nums - available_nums

    return {
        "answer": response.content,
        "sources": [
            {
                "citation": f"[{i+1}]",
                "source": chunks[i].metadata.get("source"),
                "page": chunks[i].metadata.get("page"),
                "section": chunks[i].metadata.get("h2", ""),
                "preview": chunks[i].page_content[:200] + "..."
            }
            for i in range(len(chunks))
            if (i + 1) in cited_nums
        ],
        "citation_integrity": len(hallucinated_citations) == 0,
        "hallucinated_refs": list(hallucinated_citations),
    }

from langchain.text_splitter import (
    MarkdownHeaderTextSplitter,
    RecursiveCharacterTextSplitter,
)
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.retrievers.multi_vector import MultiVectorRetriever
from langchain.storage import InMemoryByteStore
from sentence_transformers import CrossEncoder
import numpy as np
import uuid

class DocumentRAGPipeline:
    def __init__(self):
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
        self.reranker = CrossEncoder("BAAI/bge-reranker-v2-m3")
        self.vectorstore = FAISS.from_documents([], self.embeddings)
        self.docstore = InMemoryByteStore()
        self.retriever = MultiVectorRetriever(
            vectorstore=self.vectorstore,
            byte_store=self.docstore,
            id_key="doc_id",
        )

    def ingest(self, markdown: str, source: str) -> int:
        """Ingest a structured Markdown document."""
        # Step 1: Section-aware splitting
        header_splitter = MarkdownHeaderTextSplitter(
            headers_to_split_on=[("#", "h1"), ("##", "h2"), ("###", "h3")]
        )
        sections = header_splitter.split_text(markdown)

        # Step 2: Sub-split + multi-representation indexing
        char_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500, chunk_overlap=50
        )
        count = 0
        for section in sections:
            sub_chunks = char_splitter.split_text(section.page_content)
            for text in sub_chunks:
                doc_id = str(uuid.uuid4())
                original = Document(
                    page_content=text,
                    metadata={**section.metadata, "source": source, "doc_id": doc_id}
                )

                # Store original
                self.docstore.mset([(doc_id, original)])

                # Rep 1: raw text
                self.vectorstore.add_documents([original])

                # Rep 2: hypothetical questions
                qs = self.llm.invoke(
                    f"Generate 2 questions this answers:\n{text}"
                ).content
                for q in qs.strip().split("\n"):
                    if q.strip():
                        self.vectorstore.add_documents([Document(
                            page_content=q.strip(),
                            metadata={"doc_id": doc_id, "type": "question"}
                        )])
                count += 1
        return count

    def query(self, question: str, k: int = 10, top_k: int = 5) -> dict:
        """Full RAG pipeline: retrieve, rerank, generate with citations."""
        # Stage 1: Multi-vector retrieval
        candidates = self.retriever.invoke(question)[:k]

        if not candidates:
            return {"answer": "No relevant documents found.", "sources": []}

        # Stage 2: Rerank
        if len(candidates) > top_k:
            pairs = [[question, doc.page_content] for doc in candidates]
            scores = self.reranker.predict(pairs)
            ranked = np.argsort(scores)[::-1][:top_k]
            candidates = [candidates[i] for i in ranked]

        # Stage 3: Generate with citations
        return generate_with_citations(question, candidates)

# Usage
pipeline = DocumentRAGPipeline()
count = pipeline.ingest(markdown_text, source="annual_report_2025.pdf")
print(f"Indexed {count} chunks with multi-representation")

result = pipeline.query("What was the total revenue in 2025?")
print(result["answer"])
for src in result["sources"]:
    print(f"  {src['citation']}: {src['source']}, p.{src['page']}")

from ragas import evaluate
from ragas.metrics import (
    context_precision,
    context_recall,
    faithfulness,
    answer_relevancy,
)
from datasets import Dataset

# Build evaluation dataset from your RAG pipeline
eval_questions = [
    "What was total revenue in 2025?",
    "How many vacation days do employees get?",
    "What are the API rate limits?",
]

eval_data = {"question": [], "answer": [], "contexts": [], "ground_truth": []}

for q, gt in zip(eval_questions, ground_truths):
    result = pipeline.query(q)
    eval_data["question"].append(q)
    eval_data["answer"].append(result["answer"])
    eval_data["contexts"].append([c.page_content for c in result["chunks"]])
    eval_data["ground_truth"].append(gt)

dataset = Dataset.from_dict(eval_data)

results = evaluate(
    dataset,
    metrics=[context_precision, context_recall, faithfulness, answer_relevancy],
)

print(results)
# {'context_precision': 0.91, 'context_recall': 0.87,
#  'faithfulness': 0.94, 'answer_relevancy': 0.89}

# Dig into failures:
df = results.to_pandas()
low_faith = df[df["faithfulness"] < 0.7]
print(f"\n{len(low_faith)} questions with faithfulness < 0.7:")
for _, row in low_faith.iterrows():
    print(f"  Q: {row['question']}")
    print(f"  A: {row['answer'][:100]}...")
    print(f"  Faithfulness: {row['faithfulness']:.2f}\n")