How to build Enterprise Hybrid RAG

import os
import psycopg2
import openai
import cohere
import tiktoken

# ── Clients ───────────────────────────────────────────────────────
openai_client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
cohere_client = cohere.Client(api_key=os.environ["COHERE_API_KEY"])

DB_CONFIG = {
    "host": os.environ.get("PG_HOST", "localhost"),
    "port": int(os.environ.get("PG_PORT", 5432)),
    "dbname": os.environ.get("PG_DB", "ragdb"),
    "user": os.environ.get("PG_USER", "postgres"),
    "password": os.environ.get("PG_PASSWORD", ""),
}

EMBEDDING_MODEL   = "text-embedding-3-large"
EMBEDDING_DIM     = 3072
RERANK_MODEL      = "rerank-v3.5"
LLM_MODEL         = "gpt-5"
CHUNK_SIZE_TOKENS = 400   # target tokens per chunk
CHUNK_OVERLAP     = 0.20  # 20% overlap

def chunk_text(text: str, max_tokens: int = CHUNK_SIZE_TOKENS,
               overlap_ratio: float = CHUNK_OVERLAP) -> list[str]:
    """
    Token-aware sliding window chunker.
    Returns a list of text chunks.
    """
    enc = tiktoken.encoding_for_model("gpt-4o")  # cl100k_base
    tokens = enc.encode(text)
    overlap = int(max_tokens * overlap_ratio)
    step    = max_tokens - overlap

    chunks = []
    start  = 0
    while start < len(tokens):
        end        = min(start + max_tokens, len(tokens))
        chunk_toks = tokens[start:end]
        chunks.append(enc.decode(chunk_toks))
        if end == len(tokens):
            break
        start += step

    return chunks

def embed(texts: list[str]) -> list[list[float]]:
    """
    Batch embed a list of texts using text-embedding-3-large.
    OpenAI supports up to 2048 inputs per call; keep batches ≤ 500
    for safety.
    """
    response = openai_client.embeddings.create(
        model=EMBEDDING_MODEL,
        input=texts,
    )
    return [item.embedding for item in response.data]


def embed_single(text: str) -> list[float]:
    return embed([text])[0]

INSERT_SQL = """
INSERT INTO document_chunks
    (document_id, chunk_number, chunk_text, token_count, embedding, metadata)
VALUES
    (%s, %s, %s, %s, %s, %s)
ON CONFLICT DO NOTHING
"""

def ingest_document(conn, document_id: int, text: str,
                    metadata: dict = None):
    """
    Chunk, embed, and insert a document's chunks into PostgreSQL.
    The search_vector column is handled by the trigger.
    """
    enc    = tiktoken.encoding_for_model("gpt-4o")
    chunks = chunk_text(text)

    # Embed all chunks in one batch call for efficiency
    embeddings = embed(chunks)

    with conn.cursor() as cur:
        for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
            token_count = len(enc.encode(chunk))
            import json
            cur.execute(INSERT_SQL, (
                document_id,
                i,
                chunk,
                token_count,
                embedding,            # psycopg2 serialises Python list → pgvector
                json.dumps(metadata or {}),
            ))
    conn.commit()
    print(f"Ingested {len(chunks)} chunks for document {document_id}")

def dense_search(conn, query_embedding: list[float],
                 k: int = 100) -> list[tuple]:
    """
    Approximate nearest neighbour search using HNSW cosine distance.
    Returns list of (id, chunk_text) tuples.

    Note: set ef_search for query-time recall tuning:
      SET hnsw.ef_search = 200;   -- higher = better recall, slower
    """
    with conn.cursor() as cur:
        # Optionally set ef_search for this session
        cur.execute("SET hnsw.ef_search = 200;")

        cur.execute("""
            SELECT
                id,
                chunk_text,
                1 - (embedding <=> %s::vector) AS cosine_similarity
            FROM document_chunks
            ORDER BY embedding <=> %s::vector
            LIMIT %s
        """, (query_embedding, query_embedding, k))

        return cur.fetchall()   # [(id, chunk_text, similarity), ...]

def lexical_search(conn, query: str, k: int = 100) -> list[tuple]:
    """
    Full-text search using PostgreSQL tsvector + ts_rank.
    ts_rank is a BM25 approximation — good enough for most workloads.

    For true BM25, replace this with a pg_bm25 / ParadeDB query.
    """
    with conn.cursor() as cur:
        cur.execute("""
            SELECT
                id,
                chunk_text,
                ts_rank(
                    search_vector,
                    plainto_tsquery('english', %s),
                    32          -- normalization: divide by doc length
                ) AS bm25_score
            FROM document_chunks
            WHERE search_vector @@ plainto_tsquery('english', %s)
            ORDER BY bm25_score DESC
            LIMIT %s
        """, (query, query, k))

        return cur.fetchall()   # [(id, chunk_text, bm25_score), ...]

from collections import defaultdict

def rrf_merge(dense_results: list[tuple],
              lexical_results: list[tuple],
              k: int = 60) -> list[tuple]:
    """
    Reciprocal Rank Fusion.

    RRF_score(d) = Σ_i  1 / (k + rank_i(d))

    k=60 is the empirically validated default from Cormack et al., 2009.

    Parameters:
        dense_results   — list of (id, chunk_text, score) from pgvector
        lexical_results — list of (id, chunk_text, score) from FTS/BM25
        k               — smoothing constant (default 60)

    Returns:
        Merged + sorted list of (id, chunk_text, rrf_score) tuples.
    """
    scores: dict[int, float]       = defaultdict(float)
    docs:   dict[int, tuple]       = {}

    # Accumulate RRF scores from dense results (1-indexed ranks)
    for rank, row in enumerate(dense_results, start=1):
        doc_id = row[0]
        docs[doc_id] = row
        scores[doc_id] += 1.0 / (k + rank)

    # Accumulate RRF scores from lexical results
    for rank, row in enumerate(lexical_results, start=1):
        doc_id = row[0]
        docs[doc_id] = row          # lexical row overwrites; same doc_id
        scores[doc_id] += 1.0 / (k + rank)

    # Sort by descending RRF score
    merged_ids = sorted(scores, key=lambda d: scores[d], reverse=True)

    return [(doc_id, docs[doc_id][1], scores[doc_id])
            for doc_id in merged_ids]

def cohere_rerank(question: str, candidates: list[tuple],
                  top_n: int = 10) -> list[tuple]:
    """
    Rerank candidates using Cohere cross-encoder.

    candidates: list of (id, chunk_text, score) from RRF
    Returns: top_n reranked (id, chunk_text, relevance_score) tuples.
    """
    if not candidates:
        return []

    texts = [row[1] for row in candidates]

    response = cohere_client.rerank(
        query=question,
        documents=texts,
        top_n=top_n,
        model=RERANK_MODEL,
        return_documents=True,
    )

    reranked = []
    for result in response.results:
        original_row   = candidates[result.index]
        relevance_score = result.relevance_score
        reranked.append((original_row[0], original_row[1], relevance_score))

    return reranked

def build_context(chunks: list[tuple]) -> str:
    """
    Assemble numbered context blocks for the LLM prompt.
    chunks: list of (id, chunk_text, score)
    """
    parts = []
    for chunk_id, chunk_text, _ in chunks:
        parts.append(f"[Chunk {chunk_id}]\n{chunk_text.strip()}")
    return "\n\n---\n\n".join(parts)

SYSTEM_PROMPT = """You are a grounded question-answering assistant.

Rules:
1. Answer ONLY using the supplied context below.
2. Do NOT use any external knowledge or prior training data.
3. Cite the chunk IDs that support each statement (e.g., [Chunk 145]).
4. If the answer is not present in the context, respond exactly:
   "I don't have enough information to answer this question."
5. Be concise and factually precise.
"""

def build_prompt(question: str, context: str) -> str:
    return f"""{SYSTEM_PROMPT}

=== CONTEXT ===
{context}

=== QUESTION ===
{question}

=== ANSWER ==="""


def generate_answer(question: str, context: str) -> str:
    response = openai_client.chat.completions.create(
        model=LLM_MODEL,
        temperature=0,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": f"Context:\n{context}\n\nQuestion:\n{question}"},
        ],
    )
    return response.choices[0].message.content

def hybrid_rag_query(question: str, conn,
                     dense_k:   int = 100,
                     lexical_k: int = 100,
                     rrf_k:     int = 60,
                     rerank_n:  int = 10) -> dict:
    """
    Full hybrid RAG query pipeline (non-LangChain).

    Returns:
        {
            "question":      str,
            "answer":        str,
            "chunks_used":   list of (id, text, score),
            "context":       str,
        }
    """
    # 1. Embed the question
    query_embedding = embed_single(question)

    # 2. Dual retrieval
    dense_results   = dense_search(conn, query_embedding, k=dense_k)
    lexical_results = lexical_search(conn, question, k=lexical_k)

    # 3. Merge with RRF
    merged = rrf_merge(dense_results, lexical_results, k=rrf_k)

    # 4. Rerank with Cohere
    top_chunks = cohere_rerank(question, merged, top_n=rerank_n)

    # 5. Build context
    context = build_context(top_chunks)

    # 6. Generate answer
    answer  = generate_answer(question, context)

    return {
        "question":    question,
        "answer":      answer,
        "chunks_used": top_chunks,
        "context":     context,
    }


# ── Usage ──────────────────────────────────────────────────────────
if __name__ == "__main__":
    conn = psycopg2.connect(**DB_CONFIG)
    result = hybrid_rag_query(
        question="What is the capital of France?",
        conn=conn,
    )
    print(result["answer"])
    conn.close()

from langchain_core.runnables import (
    RunnableLambda,
    RunnableParallel,
    RunnablePassthrough,
)
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

# LangChain-compatible LLM and embeddings
llm = ChatOpenAI(model=LLM_MODEL, temperature=0)
embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)

# ── Step 1: Query Rewrite ──────────────────────────────────────────
rewrite_prompt = ChatPromptTemplate.from_template("""
You are a search query optimizer.
Rewrite the following question to be more precise and search-friendly.
Return ONLY the rewritten query, nothing else.

Original: {question}
Rewritten:""")

query_rewriter = (
    rewrite_prompt
    | llm
    | StrOutputParser()
)


# ── Step 2: Multi-Query Expansion ─────────────────────────────────
multi_query_prompt = ChatPromptTemplate.from_template("""
Generate 3 different search query variations for the question below.
Return one query per line. No numbering or bullets.

Question: {question}
Queries:""")

multi_query_expander = (
    multi_query_prompt
    | llm
    | StrOutputParser()
    | RunnableLambda(lambda text: [q.strip() for q in text.strip().split("\n") if q.strip()])
)


# ── Step 3: Hybrid Retriever ───────────────────────────────────────
def hybrid_retrieve(question: str, conn, dense_k=100, lexical_k=100) -> list[tuple]:
    """Runs dense + lexical search for a single query."""
    q_emb    = embed_single(question)
    dense    = dense_search(conn, q_emb, k=dense_k)
    lexical  = lexical_search(conn, question, k=lexical_k)
    return dense, lexical


def multi_query_retrieve(questions: list[str], conn) -> list[tuple]:
    """
    Runs hybrid retrieval for multiple query variants,
    then performs a single RRF over all results.
    """
    all_dense   = []
    all_lexical = []

    for q in questions:
        d, l = hybrid_retrieve(q, conn)
        all_dense.extend(d)
        all_lexical.extend(l)

    # Deduplicate by chunk id before merging
    seen_dense   = {}
    seen_lexical = {}
    for row in all_dense:
        seen_dense.setdefault(row[0], row)
    for row in all_lexical:
        seen_lexical.setdefault(row[0], row)

    return rrf_merge(
        list(seen_dense.values()),
        list(seen_lexical.values()),
    )


# ── Step 4: Reranker ───────────────────────────────────────────────
def rerank_step(data: dict) -> dict:
    question    = data["question"]
    candidates  = data["candidates"]
    top_chunks  = cohere_rerank(question, candidates, top_n=10)
    return {**data, "top_chunks": top_chunks}


# ── Step 5: Context + Prompt Assembly ─────────────────────────────
RAG_PROMPT = ChatPromptTemplate.from_template("""
You are a grounded question-answering assistant.

Answer ONLY using the supplied context.
Do NOT use external knowledge.
Cite chunk IDs for every claim (e.g., [Chunk 145]).
If the answer is not in the context, say:
"I don't have enough information to answer."

=== CONTEXT ===
{context}

=== QUESTION ===
{question}

=== ANSWER ===
""")

def build_lcel_chain(conn):
    """
    Build and return the full LCEL hybrid RAG chain.
    Requires a live psycopg2 connection.
    """

    # Step 1: Rewrite query
    rewrite = RunnableLambda(
        lambda x: {"question": query_rewriter.invoke({"question": x["question"]}),
                   "original": x["question"]}
    )

    # Step 2: Expand into multiple queries
    expand = RunnableLambda(
        lambda x: {
            **x,
            "queries": [x["question"]] + multi_query_expander.invoke({"question": x["question"]})
        }
    )

    # Step 3: Retrieve candidates
    retrieve = RunnableLambda(
        lambda x: {
            **x,
            "candidates": multi_query_retrieve(x["queries"], conn)
        }
    )

    # Step 4: Rerank
    rerank = RunnableLambda(rerank_step)

    # Step 5: Build context string
    build_ctx = RunnableLambda(
        lambda x: {**x, "context": build_context(x["top_chunks"])}
    )

    # Step 6: Generate answer
    generate = RunnableLambda(
        lambda x: RAG_PROMPT.invoke({
            "question": x["original"],
            "context":  x["context"],
        })
    ) | llm | StrOutputParser()

    # ── Full chain ─────────────────────────────────────────────────
    chain = rewrite | expand | retrieve | rerank | build_ctx | generate

    return chain


# ── Conceptual summary ─────────────────────────────────────────────
#
# chain = (
#     query_rewriter
#     | multi_query_expander
#     | hybrid_retriever
#     | rrf_merger
#     | cohere_reranker
#     | context_builder
#     | rag_prompt
#     | llm
#     | StrOutputParser()
# )

# ── Build and invoke ───────────────────────────────────────────────
conn  = psycopg2.connect(**DB_CONFIG)
chain = build_lcel_chain(conn)

result = chain.invoke({"question": "What is the capital of France?"})
print(result)

# ── Streaming (LCEL supports streaming natively) ───────────────────
for token in chain.stream({"question": "Explain the water cycle."}):
    print(token, end="", flush=True)

conn.close()

SYSTEM_PROMPT = """You are a precise, grounded QA assistant.

Strict rules:
1. Use ONLY the context provided below to answer.
2. Do NOT draw on your training data or outside knowledge.
3. For every factual claim, cite supporting chunk IDs: [Chunk 145].
4. If the answer cannot be found in the context, respond exactly:
   "I don't have enough information to answer this question."
5. Never guess, infer, or extrapolate beyond what the context states.
6. Keep your answer concise and factually accurate."""

def build_context_with_metadata(chunks: list[tuple],
                                 metadata: list[dict] = None) -> str:
    """
    Build rich context with optional source metadata.
    """
    parts = []
    for i, (chunk_id, chunk_text, score) in enumerate(chunks):
        meta_str = ""
        if metadata and i < len(metadata):
            m = metadata[i]
            src  = m.get("source", "unknown")
            page = m.get("page", "")
            meta_str = f"Source: {src}" + (f", Page: {page}" if page else "")

        block = f"[Chunk {chunk_id}]"
        if meta_str:
            block += f"  ({meta_str})"
        block += f"\n{chunk_text.strip()}"
        parts.append(block)

    return "\n\n---\n\n".join(parts)

def validate_answer(answer: str, context: str) -> dict:
    """
    Basic heuristic checks on the generated answer.
    In production, use an LLM-as-judge or Cohere groundedness API.
    """
    no_info_phrase = "I don't have enough information"
    has_citations  = "[Chunk" in answer
    refused        = no_info_phrase in answer

    return {
        "has_citations":  has_citations,
        "refused":        refused,
        "needs_review":   not has_citations and not refused,
    }

# ❌ Wrong — 0-indexed
for rank, row in enumerate(dense_results):
    scores[doc_id] += 1 / (k + rank)

# ✅ Correct — 1-indexed as per Cormack et al., 2009
for rank, row in enumerate(dense_results, start=1):
    scores[doc_id] += 1.0 / (k + rank)

# ❌ Fragile — document text may not be returned
response = cohere_client.rerank(query=question, documents=texts, top_n=10, model="rerank-v3.5")

# ✅ Explicit and safe
response = cohere_client.rerank(query=question, documents=texts, top_n=10,
                                model="rerank-v3.5", return_documents=True)

from pgvector.psycopg2 import register_vector

conn = psycopg2.connect(**DB_CONFIG)
register_vector(conn)   # ← ensures correct vector type serialisation