Hypothetical Questions - Index by Questions, Not Answers

from typing import List, Dict
import json


QUESTION_GEN_PROMPT = """Generate {n} specific questions that the following text
directly and completely answers. Phrase the questions naturally, as a user would
type them into a search box.

Text:
{chunk}

Return a JSON array of {n} question strings. No preamble."""


def generate_hypothetical_questions(
    chunk: str,
    llm,
    n: int = 5,
) -> List[str]:
    """Generate n hypothetical questions a chunk could answer."""
    prompt = QUESTION_GEN_PROMPT.format(chunk=chunk, n=n)
    response = llm.invoke(prompt)
    try:
        questions = json.loads(response)
        return [q.strip() for q in questions if isinstance(q, str)]
    except json.JSONDecodeError:
        # Fallback: parse line by line
        return [
            line.strip().lstrip("0123456789.-) ")
            for line in response.splitlines()
            if line.strip() and "?" in line
        ][:n]


def index_with_hypothetical_questions(
    documents: List[Dict],
    embed_fn,
    llm,
    n_questions: int = 5,
) -> Dict:
    """
    Build two parallel indexes:
      - question_records: one entry per generated question, pointing to its source chunk
      - chunk_records: the original chunks used for generation
    """
    chunk_records = []
    question_records = []

    for doc in documents:
        chunk_id = doc["id"]
        chunk_records.append(doc)

        questions = generate_hypothetical_questions(doc["content"], llm, n=n_questions)
        for q in questions:
            question_records.append({
                "question": q,
                "chunk_id": chunk_id,
                "source": doc["source"],
            })

    # Embed all hypothetical questions in one batch
    question_texts = [r["question"] for r in question_records]
    question_embeddings = embed_fn(question_texts)   # shape: (M, dim)

    return {
        "chunk_records": chunk_records,
        "question_records": question_records,
        "question_embeddings": question_embeddings,
    }

import numpy as np


def retrieve_by_hypothetical_questions(
    query: str,
    index: Dict,
    embed_fn,
    top_k: int = 5,
) -> List[Dict]:
    """
    Match user query against hypothetical question embeddings (Q→Q alignment).
    Return the source chunks associated with the best-matching questions.
    """
    query_vec = np.array(embed_fn([query])[0])
    embeddings = np.array(index["question_embeddings"])

    scores = embeddings @ query_vec / (
        np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_vec) + 1e-9
    )
    top_indices = np.argsort(scores)[::-1]

    # Collect unique source chunks (multiple questions may point to the same chunk)
    seen_chunk_ids = set()
    results = []
    chunk_lookup = {r["id"]: r for r in index["chunk_records"]}

    for idx in top_indices:
        if len(results) >= top_k:
            break
        q_record = index["question_records"][idx]
        chunk_id = q_record["chunk_id"]
        if chunk_id not in seen_chunk_ids:
            seen_chunk_ids.add(chunk_id)
            results.append({
                "chunk": chunk_lookup[chunk_id],
                "matched_question": q_record["question"],
                "score": float(scores[idx]),
            })

    return results


def hypothetical_questions_rag(
    query: str,
    index: Dict,
    embed_fn,
    llm,
) -> Dict:
    # Retrieve via Q→Q alignment
    retrieved = retrieve_by_hypothetical_questions(query, index, embed_fn)

    # Build generation context from source chunks (not the generated questions)
    context = "\n\n".join([r["chunk"]["content"] for r in retrieved])

    response = llm.invoke(f"Context:\n{context}\n\nQuestion: {query}")

    return {
        "answer": response,
        "retrieved": retrieved,
        "matched_questions": [r["matched_question"] for r in retrieved],
    }

from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct, VectorParams, Distance


def build_qdrant_index(index: Dict, collection_name: str = "hyp_questions"):
    """Store hypothetical question vectors in Qdrant with chunk payload."""
    client = QdrantClient(":memory:")
    dim = len(index["question_embeddings"][0])

    client.create_collection(
        collection_name=collection_name,
        vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
    )

    points = [
        PointStruct(
            id=i,
            vector=index["question_embeddings"][i],
            payload={
                "question": index["question_records"][i]["question"],
                "chunk_id": index["question_records"][i]["chunk_id"],
                "source": index["question_records"][i]["source"],
            },
        )
        for i in range(len(index["question_records"]))
    ]
    client.upsert(collection_name=collection_name, points=points)
    return client