from langchain_text_splitters import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=1200,       # Characters — not tokens
    chunk_overlap=200,     # 15-20% overlap prevents information loss at boundaries
    separators=[
        "\n## ",           # H2 headers (highest priority split point)
        "\n### ",          # H3 headers
        "\n\n",            # Paragraph breaks
        "\n",              # Line breaks
        ". ",              # Sentence boundaries
        " ",               # Word boundaries (last resort)
    ],
    length_function=len,
    is_separator_regex=False,
)

chunks = splitter.split_documents(raw_docs)

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.storage import InMemoryStore
from langchain.retrievers import ParentDocumentRetriever

# Parent chunks: large context windows for generation
parent_splitter = RecursiveCharacterTextSplitter(chunk_size=4000, chunk_overlap=400)

# Child chunks: small, precise for retrieval matching
child_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)

store = InMemoryStore()  # Use Redis or PostgreSQL in production

retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_splitter,
    parent_splitter=parent_splitter,
)

-- Enable extension
CREATE EXTENSION IF NOT EXISTS vector;

-- Create embeddings table with HNSW index
CREATE TABLE document_embeddings (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    content TEXT NOT NULL,
    metadata JSONB NOT NULL DEFAULT '{}',
    embedding vector(1024) NOT NULL,  -- voyage-3 dimensions
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- HNSW index for approximate nearest neighbor search
CREATE INDEX ON document_embeddings
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 128);

-- GIN index on metadata for filtered retrieval
CREATE INDEX ON document_embeddings
USING gin (metadata jsonb_path_ops);

from langchain_postgres import PGVector
from langchain_voyageai import VoyageAIEmbeddings

embeddings = VoyageAIEmbeddings(
    model="voyage-3",
    voyage_api_key=os.environ["VOYAGE_API_KEY"],
)

vectorstore = PGVector(
    embeddings=embeddings,
    collection_name="compliance_docs",
    connection=os.environ["DATABASE_URL"],  # postgresql://...
    use_jsonb=True,
)

# Index documents
vectorstore.add_documents(chunks)

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

bm25_retriever = BM25Retriever.from_documents(chunks, k=10)
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 10})

hybrid_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.3, 0.7],  # 70% semantic, 30% keyword
)

from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank

reranker = CohereRerank(
    model="rerank-english-v3.0",
    cohere_api_key=os.environ["COHERE_API_KEY"],
    top_n=5,
)

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

# Retrieve + rerank
results = compression_retriever.invoke("What are the data retention requirements for HIPAA?")

from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate

llm = ChatAnthropic(
    model="claude-sonnet-4-20250514",
    temperature=0,
    max_tokens=4096,
    anthropic_api_key=os.environ["ANTHROPIC_API_KEY"],
)

prompt = ChatPromptTemplate.from_messages([
    ("system", """You are a compliance knowledge assistant for an enterprise organization.
Answer questions using ONLY the provided context documents. For each claim in your answer,
cite the source using [Source: document_name, page X] format.

If the context does not contain sufficient information to answer the question,
say "I don't have enough information in the available documents to answer this question"
and suggest what additional documents might help.

Do not use prior knowledge. Do not speculate. Be precise and specific."""),
    ("human", """Context documents:
{context}

Question: {question}

Answer with citations:"""),
])

chain = prompt | llm

from langchain_core.runnables import RunnablePassthrough

rag_chain = (
    {
        "context": compression_retriever | format_docs,
        "question": RunnablePassthrough(),
    }
    | prompt
    | llm
)

# Stream response
async for chunk in rag_chain.astream("What are HIPAA data retention requirements?"):
    print(chunk.content, end="", flush=True)

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

# Prepare evaluation dataset
eval_results = evaluate(
    dataset=eval_dataset,
    metrics=[faithfulness, answer_relevancy, context_precision, context_recall],
    llm=evaluation_llm,
    embeddings=embeddings,
)

print(eval_results)
# {'faithfulness': 0.92, 'answer_relevancy': 0.88,
#  'context_precision': 0.85, 'context_recall': 0.79}