full_inference.py

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

from imports import HuggingFaceEmbeddings, LangchainEmbedding, Settings, genai
from pipeline import TableSimilarityCache
from run_base_retrieval import (
    load_index as base_load_index,
    load_questions as base_load_questions,
    run_retrieval as base_run_retrieval,
)
from run_table_expansion import (
    build_faiss_store,
    expand_question,
)
from run_table_pruning import prune_question
from run_sql_generation import (
    SYSTEM_PROMPT,
    call_openai_chat,
    cleanup_sql,
    estimate_cost,
    estimate_tokens,
    format_tables_for_prompt,
    build_prompt,
)
from imports import CrossEncoder


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_env_file(path: Optional[str]) -> None:
    if not path:
        return
    env_path = Path(path)
    if not env_path.exists():
        raise FileNotFoundError(f"Environment file '{path}' not found.")

    with env_path.open("r", encoding="utf-8") as handle:
        for raw_line in handle:
            line = raw_line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, value = line.split("=", 1)
            key = key.strip()
            value = value.strip()
            if key and value and key not in os.environ:
                os.environ[key] = value


def save_json(path: Path, payload: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2, ensure_ascii=False)
    logging.info("Saved %s", path)


def stage_base_retrieval(args: argparse.Namespace) -> Dict[str, Any]:
    questions = base_load_questions(args.questions_file, args.question_key, args.id_key)
    index = base_load_index(args.vector_store_path, args.embedding_model, args.chunk_size)
    retrievals = base_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,
    }
    logging.info("Stage 1 complete: retrieved tables for %d questions", len(retrievals))
    return summary


def stage_expansion(args: argparse.Namespace, retrieval_summary: Dict[str, Any]) -> Dict[str, Any]:
    logging.info("Loading embedding model %s for joinable discovery", args.join_embedding_model)
    lc_embed_model = HuggingFaceEmbeddings(model_name=args.join_embedding_model)
    Settings.embed_model = LangchainEmbedding(lc_embed_model)

    faiss_store = build_faiss_store(args.faiss_index, args.faiss_metadata)

    expanded_questions: List[Dict[str, Any]] = []
    start = time.time()
    for question in retrieval_summary["questions"]:
        expanded_questions.append(
            expand_question(
                question,
                lc_embed_model,
                faiss_store,
                per_table_limit=args.join_per_table,
                total_limit=args.join_total_limit,
                threshold=args.join_threshold,
            )
        )
    elapsed = time.time() - start

    summary = {
        "retrieval_metadata": {
            "questions_file": retrieval_summary.get("questions_file"),
            "vector_store_path": retrieval_summary.get("vector_store_path"),
        },
        "faiss_index": args.faiss_index,
        "faiss_metadata": args.faiss_metadata,
        "embedding_model": args.join_embedding_model,
        "per_table_limit": args.join_per_table,
        "total_limit": args.join_total_limit,
        "threshold": args.join_threshold,
        "expanded_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "processing_time_sec": elapsed,
        "questions": expanded_questions,
    }
    logging.info("Stage 2 complete: expanded tables for %d questions", len(expanded_questions))
    return summary


def stage_pruning(args: argparse.Namespace, expansion_summary: Dict[str, Any]) -> Dict[str, Any]:
    logging.info("Loading table repository from %s", args.table_repository)
    with open(args.table_repository, "r", encoding="utf-8") as handle:
        table_repository = json.load(handle)

    logging.info("Loading cross-encoder model %s", args.cross_encoder_model)
    cross_encoder = CrossEncoder(
        args.cross_encoder_model,
        automodel_args={"torch_dtype": "auto"},
        trust_remote_code=True,
    )

    cache = TableSimilarityCache(cache_dir=args.cache_dir, cache_version=args.cache_version)

    pruned_questions: List[Dict[str, Any]] = []
    start = time.time()
    for question in expansion_summary["questions"]:
        pruned_questions.append(
            prune_question(
                question,
                table_repository,
                cross_encoder,
                cache,
                alpha=args.pruning_alpha,
                beta=args.pruning_beta,
                top_n=args.pruning_top_n,
                max_workers=args.pruning_max_workers,
            )
        )
    elapsed = time.time() - start

    summary = {
        "table_repository": args.table_repository,
        "cross_encoder_model": args.cross_encoder_model,
        "alpha": args.pruning_alpha,
        "beta": args.pruning_beta,
        "top_n": args.pruning_top_n,
        "max_workers": args.pruning_max_workers,
        "cache_dir": args.cache_dir,
        "cache_version": args.cache_version,
        "pruned_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "processing_time_sec": elapsed,
        "questions": pruned_questions,
    }
    logging.info("Stage 3 complete: pruned table sets for %d questions", len(pruned_questions))
    return summary, table_repository


def stage_sql_generation(args: argparse.Namespace, pruning_summary: Dict[str, Any], table_repository: Dict[str, Any]) -> Dict[str, Any]:
    provider = args.provider.lower()

    if provider == "gemini":
        api_key = args.api_key or os.getenv("GOOGLE_API_KEY")
        if not api_key:
            raise ValueError("Gemini provider selected but no API key supplied (use --api-key or GOOGLE_API_KEY).")
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel(
            args.model_name,
            generation_config={
                "temperature": args.temperature,
                "max_output_tokens": args.max_output_tokens,
            },
        )
        client_info: Dict[str, Any] = {"provider": "gemini", "client": model}
    elif provider in {"openai", "deepinfra"}:
        if provider == "openai":
            default_base = "https://api.openai.com/v1"
            default_env = "OPENAI_API_KEY"
            if args.model_name == "models/gemini-2.0-flash":
                args.model_name = "gpt-4o-mini"
        else:
            default_base = "https://api.deepinfra.com/v1/openai"
            default_env = "DEEPINFRA_API_KEY"
            if args.model_name == "models/gemini-2.0-flash":
                args.model_name = "meta-llama/Llama-3.2-3B-Instruct"

        api_key = args.api_key or os.getenv(default_env)
        if not api_key:
            raise ValueError(
                f"{provider} provider selected but no API key supplied (use --api-key or {default_env})."
            )
        api_base = args.api_base or default_base
        client_info = {
            "provider": provider,
            "api_key": api_key,
            "api_base": api_base,
            "organization": args.organization,
        }
    else:
        raise ValueError(f"Unsupported provider '{args.provider}'.")

    results: List[Dict[str, Any]] = []
    total_input_tokens = 0
    total_output_tokens = 0
    total_cost = 0.0

    for entry in pruning_summary["questions"]:
        question_id = entry.get("question_id")
        question_text = entry.get("question", "")
        table_ids = entry.get("pruned_table_ids") or []

        if not question_text or not table_ids:
            results.append(
                {
                    "question_id": question_id,
                    "question": question_text,
                    "table_ids": table_ids,
                    "sql": "",
                    "provider": provider,
                    "error": "Missing question text or table IDs.",
                }
            )
            continue

        tables_info = format_tables_for_prompt(table_repository, table_ids, args.max_sample_rows)
        if not tables_info:
            results.append(
                {
                    "question_id": question_id,
                    "question": question_text,
                    "table_ids": table_ids,
                    "sql": "",
                    "provider": provider,
                    "error": "Referenced tables not found in repository.",
                }
            )
            continue

        prompt = build_prompt(question_text, tables_info)

        try:
            if provider == "gemini":
                response = client_info["client"].generate_content(prompt)
                raw_sql = response.text if hasattr(response, "text") else str(response)
            else:
                raw_sql = call_openai_chat(
                    api_base=client_info["api_base"],
                    api_key=client_info["api_key"],
                    model=args.model_name,
                    system_prompt=SYSTEM_PROMPT,
                    user_prompt=prompt,
                    temperature=args.temperature,
                    max_tokens=args.max_output_tokens,
                    organization=client_info.get("organization"),
                )

            sql_text = cleanup_sql(raw_sql)
            input_tokens = estimate_tokens(prompt)
            output_tokens = estimate_tokens(sql_text)
            costs = estimate_cost(input_tokens, output_tokens, provider)

            total_input_tokens += input_tokens
            total_output_tokens += output_tokens
            total_cost += costs["total_cost"]

            results.append(
                {
                    "question_id": question_id,
                    "question": question_text,
                    "table_ids": table_ids,
                    "provider": provider,
                    "sql": sql_text,
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "costs": costs,
                }
            )
        except Exception as exc:
            logging.exception("SQL generation failed for question %s", question_id)
            results.append(
                {
                    "question_id": question_id,
                    "question": question_text,
                    "table_ids": table_ids,
                    "provider": provider,
                    "sql": "",
                    "error": str(exc),
                }
            )

    summary = {
        "provider": provider,
        "model_name": args.model_name,
        "temperature": args.temperature,
        "max_output_tokens": args.max_output_tokens,
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "totals": {
            "input_tokens": total_input_tokens,
            "output_tokens": total_output_tokens,
            "estimated_cost": total_cost,
            "questions_processed": len(results),
        },
        "questions": results,
    }
    logging.info("Stage 4 complete: generated SQL for %d questions", len(results))
    return summary


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run the full retrieval → expansion → pruning → SQL generation pipeline."
    )

    parser.add_argument(
        "--env-file",
        default=None,
        help="Optional path to a .env file containing API credentials (KEY=VALUE per line).",
    )

    # Stage 1: Retrieval
    parser.add_argument("--questions-file", required=True, help="Path to questions JSON.")
    parser.add_argument("--vector-store-path", required=True, help="Persisted LlamaIndex/FAISS directory.")
    parser.add_argument("--embedding-model", default="BAAI/bge-base-en-v1.5", help="Retrieval embedding model.")
    parser.add_argument("--chunk-size", type=int, default=7500, help="LlamaIndex chunk size.")
    parser.add_argument("--top-k", type=int, default=5, help="Tables to retrieve per question.")
    parser.add_argument("--question-key", default="question", help="Key holding question text.")
    parser.add_argument("--id-key", default=None, help="Optional key holding question id.")

    # Stage 2: Expansion
    parser.add_argument("--faiss-index", required=True, help="Column-level FAISS index path.")
    parser.add_argument("--faiss-metadata", required=True, help="Column metadata JSON path.")
    parser.add_argument("--join-embedding-model", default="BAAI/bge-base-en-v1.5", help="Embedding model for join discovery.")
    parser.add_argument("--join-per-table", type=int, default=3, help="Joinable tables per retrieved table.")
    parser.add_argument("--join-total-limit", type=int, default=3, help="Total joinable tables to add per question.")
    parser.add_argument("--join-threshold", type=float, default=0.7, help="Distance threshold for joinability.")

    # Stage 3: Pruning
    parser.add_argument("--table-repository", required=True, help="Schema + samples JSON.")
    parser.add_argument("--cross-encoder-model", default="jinaai/jina-reranker-v2-base-multilingual", help="Cross-encoder ID for scoring.")
    parser.add_argument("--pruning-alpha", type=float, default=0.7, help="Weight for query-table scores.")
    parser.add_argument("--pruning-beta", type=float, default=0.3, help="Weight for table-table scores.")
    parser.add_argument("--pruning-top-n", type=int, default=5, help="Tables to keep after pruning.")
    parser.add_argument("--pruning-max-workers", type=int, default=2, help="Threads for table-table scoring.")
    parser.add_argument("--cache-dir", default="table_cache", help="Directory for similarity cache.")
    parser.add_argument("--cache-version", default="v1", help="Cache version key.")

    # Stage 4: SQL generation
    parser.add_argument("--provider", choices=["gemini", "openai", "deepinfra"], default="gemini", help="LLM provider.")
    parser.add_argument("--model-name", default="models/gemini-2.0-flash", help="Model identifier to use.")
    parser.add_argument("--api-key", default=None, help="API key for provider (falls back to env).")
    parser.add_argument("--api-base", default=None, help="Override base URL for OpenAI-compatible providers.")
    parser.add_argument("--organization", default=None, help="Optional OpenAI organization header.")
    parser.add_argument("--temperature", type=float, default=0.0, help="Sampling temperature.")
    parser.add_argument("--max-output-tokens", type=int, default=1024, help="Max tokens for SQL generation.")
    parser.add_argument("--max-sample-rows", type=int, default=3, help="Rows per table when formatting prompts.")

    # Output / logging
    parser.add_argument("--output-file", required=True, help="Path to final aggregated results JSON.")
    parser.add_argument("--intermediate-dir", default=None, help="Optional directory to also store stage outputs.")
    parser.add_argument("--log-level", default="INFO", help="Logging level (DEBUG, INFO, ...).")

    return parser.parse_args()


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

    retrieval_summary = stage_base_retrieval(args)
    if args.intermediate_dir:
        save_json(Path(args.intermediate_dir) / "stage1_retrieval.json", retrieval_summary)

    expansion_summary = stage_expansion(args, retrieval_summary)
    if args.intermediate_dir:
        save_json(Path(args.intermediate_dir) / "stage2_expansion.json", expansion_summary)

    pruning_summary, table_repository = stage_pruning(args, expansion_summary)
    if args.intermediate_dir:
        save_json(Path(args.intermediate_dir) / "stage3_pruning.json", pruning_summary)

    sql_summary = stage_sql_generation(args, pruning_summary, table_repository)
    if args.intermediate_dir:
        save_json(Path(args.intermediate_dir) / "stage4_sql.json", sql_summary)

    final_payload = {
        "retrieval": retrieval_summary,
        "expansion": expansion_summary,
        "pruning": pruning_summary,
        "sql_generation": sql_summary,
    }
    save_json(Path(args.output_file), final_payload)

    print("End-to-end pipeline complete.")
    print("------------------------------")
    print(f"Questions processed: {len(sql_summary['questions'])}")
    totals = sql_summary.get("totals", {})
    print(f"Total input tokens (est.): {totals.get('input_tokens', 0)}")
    print(f"Total output tokens (est.): {totals.get('output_tokens', 0)}")
    print(f"Total estimated cost: ${totals.get('estimated_cost', 0.0):.4f}")


if __name__ == "__main__":
    main()