generate.py

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

import requests

from imports import genai


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("#"):
                continue
            if "=" 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 load_json(path: str) -> Any:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def normalize_path(path: Optional[str]) -> List[str]:
    if not path:
        return []
    return [segment for segment in path.split(".") if segment]


def extract_field(data: Dict[str, Any], path: Sequence[str], fallback_keys: Sequence[str]) -> Any:
    if path:
        value = data
        for key in path:
            if isinstance(value, dict):
                value = value.get(key)
            else:
                value = None
                break
        if value is not None:
            return value
    for key in fallback_keys:
        if key in data and data[key] is not None:
            return data[key]
    return None


def format_tables_for_prompt(
    table_repository: Dict[str, Any],
    table_ids: Sequence[str],
    max_rows: int,
) -> str:
    selected_ids = set(table_ids)
    if not selected_ids:
        return ""

    formatted_blocks: List[str] = []
    for table_name, variants in table_repository.get("tables", {}).items():
        for variant in variants:
            table_id = variant.get("id")
            if table_id not in selected_ids:
                continue

            columns = variant.get("columns", [])
            data = variant.get("data", variant.get("content", [])) or []
            description = variant.get("table_description", "No description available")

            block_lines = [
                f"=== TABLE: {table_name} (ID: {table_id}) ===",
                f"DESCRIPTION: {description}",
                "",
                f"COLUMNS: {', '.join(columns)}" if columns else "COLUMNS: (none documented)",
                "",
            ]

            if data:
                block_lines.append("SAMPLE DATA:")
                header_row = " | ".join(columns)
                block_lines.append(header_row)
                block_lines.append("-" * len(header_row))
                for row in data[:max_rows]:
                    row_values = [str(row[i]) if i < len(row) and row[i] is not None else "NULL" for i in range(len(columns))]
                    block_lines.append(" | ".join(row_values))
                block_lines.append("")

            formatted_blocks.append("\n".join(block_lines))

    return "\n".join(formatted_blocks)


def cleanup_sql(raw_text: str) -> str:
    sql_query = raw_text.strip()
    # Remove markdown fences if present
    sql_query = re.sub(r"^```sql\s*", "", sql_query, flags=re.IGNORECASE | re.MULTILINE)
    sql_query = re.sub(r"^```\s*", "", sql_query, flags=re.MULTILINE)
    sql_query = re.sub(r"\s*```$", "", sql_query, flags=re.MULTILINE)
    return sql_query.strip()


def estimate_tokens(text: str) -> int:
    return max(1, len(text) // 4)


def estimate_cost(input_tokens: int, output_tokens: int, provider: str) -> Dict[str, float]:
    provider = provider.lower()

    if provider == "gemini":
        # Gemini 2.0 Flash (Oct 2024) pricing assumption: $0.075 / 1M input tokens, $0.30 / 1M output tokens (128k context)
        input_rate = 0.075
        output_rate = 0.30
    elif provider == "openai":
        # Default to GPT-4o-mini pricing (Oct 2024): $5.00 / 1M input, $15.00 / 1M output
        input_rate = 5.00
        output_rate = 15.00
    elif provider == "deepinfra":
        # DeepInfra pricing varies by model; assume $0.80 / 1M tokens for both directions as a placeholder
        input_rate = 0.80
        output_rate = 0.80
    else:
        input_rate = 0.0
        output_rate = 0.0

    input_cost = (input_tokens / 1_000_000) * input_rate
    output_cost = (output_tokens / 1_000_000) * output_rate
    total_cost = input_cost + output_cost
    return {"input_cost": input_cost, "output_cost": output_cost, "total_cost": total_cost}


SYSTEM_PROMPT = (
    "You are an expert SQL query generator. Generate precise SQL that answers the user's question "
    "using only the provided schema. Return ONLY the SQL query without explanation."
)


def build_prompt(question: str, tables_info: str) -> str:
    return f"""You are an expert SQL query generator. Given the following tables and a natural language question, generate a precise SQL query that answers the question.

AVAILABLE TABLES:
{tables_info}

QUESTION: {question}

INSTRUCTIONS:
1. Analyze the question carefully to understand what information is being requested.
2. Identify which tables and columns are needed from the available tables.
3. Generate a syntactically correct SQL query that answers the question.
4. Use appropriate JOINs if multiple tables are needed.
5. Apply proper filtering, grouping, and ordering as required.
6. Use the exact column names and table names as shown in the schema.
7. Be careful with column names that contain spaces - use backticks or quotes as needed.
8. Return ONLY the SQL query without any explanation or markdown formatting.

SQL QUERY:"""


def call_openai_chat(
    api_base: str,
    api_key: str,
    model: str,
    system_prompt: str,
    user_prompt: str,
    temperature: float,
    max_tokens: int,
    organization: Optional[str] = None,
) -> str:
    url = api_base.rstrip("/") + "/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    if organization:
        headers["OpenAI-Organization"] = organization

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        "temperature": temperature,
        "max_tokens": max_tokens,
    }

    response = requests.post(url, headers=headers, json=payload, timeout=120)
    response.raise_for_status()
    data = response.json()

    try:
        return data["choices"][0]["message"]["content"]
    except (KeyError, IndexError) as exc:
        raise RuntimeError(f"Unexpected OpenAI/DeepInfra response: {data}") from exc


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate SQL for pruned table sets using Gemini."
    )
    parser.add_argument(
        "--env-file",
        default=None,
        help="Optional path to a .env file containing API credentials (KEY=VALUE per line).",
    )
    parser.add_argument(
        "--provider",
        choices=["gemini", "openai", "deepinfra"],
        default="gemini",
        help="LLM provider to use for SQL generation.",
    )
    parser.add_argument(
        "--pruned-file",
        required=True,
        help="JSON output produced by run_table_pruning.py (or similar).",
    )
    parser.add_argument(
        "--table-repository",
        required=True,
        help="Path to the table repository JSON used during retrieval.",
    )
    parser.add_argument(
        "--output-file",
        required=True,
        help="Destination JSON file for SQL generation results.",
    )
    parser.add_argument(
        "--question-id-field",
        default="question_id",
        help="Dot-delimited path to the question identifier in the pruned file.",
    )
    parser.add_argument(
        "--table-ids-field",
        default="pruned_table_ids",
        help="Field storing the list of table IDs to use for SQL generation.",
    )
    parser.add_argument(
        "--question-field",
        default="question",
        help="Dot-delimited path to the natural language question.",
    )
    parser.add_argument(
        "--max-sample-rows",
        type=int,
        default=3,
        help="Number of sample rows per table when formatting the prompt.",
    )
    parser.add_argument(
        "--model-name",
        default="models/gemini-2.0-flash",
        help="Model identifier to use for SQL generation (defaults vary by provider).",
    )
    parser.add_argument(
        "--api-key",
        default=None,
        help="API key for the selected provider (falls back to provider-specific environment variables).",
    )
    parser.add_argument(
        "--api-base",
        default=None,
        help="Base URL for OpenAI-compatible providers (default varies by provider).",
    )
    parser.add_argument(
        "--organization",
        default=None,
        help="Optional OpenAI organization header.",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Sampling temperature passed to Gemini.",
    )
    parser.add_argument(
        "--max-output-tokens",
        type=int,
        default=1024,
        help="Maximum tokens Gemini should return for each SQL generation request.",
    )
    parser.add_argument(
        "--log-level",
        default="INFO",
        help="Logging level (DEBUG, INFO, WARNING, ERROR).",
    )
    parser.add_argument(
        "--include-details",
        action="store_true",
        help="Include per-question SQL outputs in stdout summary.",
    )
    return parser.parse_args()


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

    provider = args.provider.lower()

    if provider == "gemini":
        default_key_env = "GOOGLE_API_KEY"
        api_key = args.api_key or os.getenv(default_key_env)
        if not api_key:
            raise ValueError(
                "A Google Gemini API key must be provided via --api-key or GOOGLE_API_KEY environment variable."
            )
        genai.configure(api_key=api_key)
        gemini_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": gemini_model,
            "api_key": api_key,
        }
    elif provider in {"openai", "deepinfra"}:
        if provider == "openai":
            default_base = "https://api.openai.com/v1"
            default_key_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_key_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_key_env)
        if not api_key:
            raise ValueError(
                f"An API key must be provided via --api-key or {default_key_env} for provider '{provider}'."
            )
        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 '{provider}'.")

    pruned_payload = load_json(args.pruned_file)
    table_repository = load_json(args.table_repository)

    questions = pruned_payload["questions"] if isinstance(pruned_payload, dict) else pruned_payload

    question_id_path = normalize_path(args.question_id_field)
    question_field_path = normalize_path(args.question_field)
    results: List[Dict[str, Any]] = []

    total_input_tokens = 0
    total_output_tokens = 0
    total_cost = 0.0

    for entry in questions:
        question_id = extract_field(entry, question_id_path, ("question_id", "id", "questionId"))
        question_text = extract_field(entry, question_field_path, ("question", "query", "text"))
        table_ids = entry.get(args.table_ids_field) or []

        if not question_text:
            logging.warning("Skipping entry without question text (question_id=%s)", question_id)
            continue

        if not table_ids:
            logging.info(
                "Skipping SQL generation for question %s (no table IDs in field '%s').",
                question_id,
                args.table_ids_field,
            )
            results.append(
                {
                    "question_id": question_id,
                    "question": question_text,
                    "table_ids": table_ids,
                    "sql": "",
                    "error": "No tables available for SQL generation.",
                }
            )
            continue

        tables_info = format_tables_for_prompt(table_repository, table_ids, args.max_sample_rows)
        if not tables_info:
            logging.warning(
                "Question %s references tables missing from repository; skipping SQL generation.",
                question_id,
            )
            results.append(
                {
                    "question_id": question_id,
                    "question": question_text,
                    "table_ids": table_ids,
                    "sql": "",
                    "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,
                    "prompt": prompt,
                    "sql": sql_text,
                    "provider": provider,
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "costs": costs,
                }
            )
            logging.info(
                "Generated SQL for question %s (tables=%d, tokens=%d/%d).",
                question_id,
                len(table_ids),
                input_tokens,
                output_tokens,
            )

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

    summary = {
        "pruned_file": args.pruned_file,
        "table_repository": args.table_repository,
        "model_name": args.model_name,
        "provider": provider,
        "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,
    }

    os.makedirs(os.path.dirname(args.output_file) or ".", exist_ok=True)
    with open(args.output_file, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)
    logging.info("Saved SQL generation results to %s", args.output_file)

    print("SQL Generation Summary")
    print("----------------------")
    print(f"Questions processed: {len(results)}")
    print(f"Total input tokens (est.): {total_input_tokens}")
    print(f"Total output tokens (est.): {total_output_tokens}")
    print(f"Total estimated cost: ${total_cost:.4f}")

    if args.include_details:
        print("\nPer-question SQL:")
        for item in results:
            sql = item.get("sql")
            if sql:
                print(f"- Q{item.get('question_id')}: {sql}")
            else:
                print(f"- Q{item.get('question_id')}: [ERROR] {item.get('error')}")


if __name__ == "__main__":
    main()