retrieve.py

import argparse
import json
import logging
import os
import sys
import time
from typing import Any, Dict, List, Optional

from imports import (
    HuggingFaceEmbeddings,
    LangchainEmbedding,
    Settings,
    StorageContext,
    load_index_from_storage,
    FaissVectorStore,
)


def configure_logging(level: str) -> None:
    logging.basicConfig(
        stream=sys.stdout,
        level=getattr(logging, level.upper(), logging.INFO),
        format="%(asctime)s %(levelname)s %(name)s - %(message)s",
    )


def load_questions(path: str, question_key: str, id_key: Optional[str]) -> List[Dict[str, Any]]:
    with open(path, "r", encoding="utf-8") as f:
        payload = json.load(f)

    if isinstance(payload, dict) and "questions" in payload:
        payload = payload["questions"]

    questions: List[Dict[str, Any]] = []
    for idx, entry in enumerate(payload):
        if isinstance(entry, str):
            question_text = entry
            question_id = idx
        elif isinstance(entry, dict):
            if question_key not in entry:
                # fall back to common alternatives
                for fallback in ("question", "query", "text"):
                    if fallback in entry:
                        question_text = entry[fallback]
                        break
                else:
                    raise KeyError(f"Question key '{question_key}' not found in entry {idx}")
            else:
                question_text = entry[question_key]

            if id_key and id_key in entry:
                question_id = entry[id_key]
            else:
                question_id = entry.get("question_id") or entry.get("id") or idx
        else:
            raise TypeError(f"Unsupported question entry type: {type(entry)} at index {idx}")

        questions.append(
            {
                "id": question_id,
                "question": question_text,
                "raw": entry,
            }
        )

    return questions


def to_serializable(value: Any) -> Any:
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, list):
        return [to_serializable(item) for item in value]
    if isinstance(value, dict):
        return {str(key): to_serializable(val) for key, val in value.items()}

    try:
        import numpy as np

        if isinstance(value, np.generic):
            return value.item()
        if isinstance(value, np.ndarray):
            return value.tolist()
    except ImportError:
        pass

    return str(value)


def extract_document_text(document: Any) -> str:
    for attr in ("text", "get_content"):
        if hasattr(document, attr):
            candidate = getattr(document, attr)
            if callable(candidate):
                try:
                    return candidate()
                except TypeError:
                    continue
            return candidate

    if hasattr(document, "text_resource"):
        text_res = getattr(document, "text_resource")
        if hasattr(text_res, "text"):
            return text_res.text

    return ""


def load_index(vector_store_path: str, embedding_model_name: str, chunk_size: int) -> Any:
    logging.info("Loading embedding model %s", embedding_model_name)
    lc_embed_model = HuggingFaceEmbeddings(model_name=embedding_model_name)
    Settings.embed_model = LangchainEmbedding(lc_embed_model)
    Settings.chunk_size = chunk_size

    logging.info("Loading persisted vector store from %s", vector_store_path)
    vector_store = FaissVectorStore.from_persist_dir(vector_store_path)
    storage_context = StorageContext.from_defaults(
        vector_store=vector_store, persist_dir=vector_store_path
    )
    index = load_index_from_storage(storage_context=storage_context)
    return index


def run_retrieval(index: Any, questions: List[Dict[str, Any]], top_k: int) -> List[Dict[str, Any]]:
    retriever = index.as_retriever(similarity_top_k=top_k)
    results: List[Dict[str, Any]] = []

    for position, item in enumerate(questions):
        question_text = item["question"]
        question_id = item["id"]
        logging.info("Retrieving tables for question %s (%s)", question_id, question_text[:80])

        start = time.time()
        retrieved_nodes = retriever.retrieve(question_text)
        elapsed = time.time() - start

        serialized_nodes = []
        for node in retrieved_nodes:
            doc = getattr(node, "node", node)
            metadata = getattr(doc, "metadata", {}) or {}
            serialized_nodes.append(
                {
                    "table_id": metadata.get("table_id"),
                    "table_name": metadata.get("table_name"),
                    "score": float(getattr(node, "score", 0.0) or 0.0),
                    "metadata": to_serializable(metadata),
                    "text": extract_document_text(doc),
                }
            )

        results.append(
            {
                "question_id": question_id,
                "question": question_text,
                "position": position,
                "retrieval_time_sec": elapsed,
                "results": serialized_nodes,
            }
        )

    return results


def save_results(output_path: str, payload: Dict[str, Any]) -> None:
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)
    logging.info("Saved retrieval results to %s", output_path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run base table retrieval and export candidate tables per question."
    )
    parser.add_argument("--questions-file", required=True, help="Path to questions JSON file.")
    parser.add_argument(
        "--vector-store-path",
        required=True,
        help="Directory containing the persisted LlamaIndex/FAISS vector store.",
    )
    parser.add_argument(
        "--embedding-model",
        default="BAAI/bge-base-en-v1.5",
        help="HuggingFace embedding model used when the vector store was created.",
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=7500,
        help="Chunk size to register with the LlamaIndex Settings. Should match persistence.",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=5,
        help="Number of tables to retrieve per question.",
    )
    parser.add_argument(
        "--output-file",
        required=True,
        help="Destination JSON file for retrieval results.",
    )
    parser.add_argument(
        "--question-key",
        default="question",
        help="Key inside each question object that holds the natural language question.",
    )
    parser.add_argument(
        "--id-key",
        default=None,
        help="Optional key holding the question identifier. Falls back to 'question_id', then index.",
    )
    parser.add_argument(
        "--log-level",
        default="INFO",
        help="Logging level (DEBUG, INFO, WARNING, ERROR).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    configure_logging(args.log_level)

    questions = load_questions(args.questions_file, args.question_key, args.id_key)
    index = load_index(args.vector_store_path, args.embedding_model, args.chunk_size)
    retrievals = run_retrieval(index, questions, args.top_k)

    summary = {
        "questions_file": args.questions_file,
        "vector_store_path": args.vector_store_path,
        "embedding_model": args.embedding_model,
        "top_k": args.top_k,
        "retrieved_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "questions": retrievals,
    }
    save_results(args.output_file, summary)


if __name__ == "__main__":
    main()