src/data/VariationGenerator.py

import json
import re
import time
from typing import List

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI


class VariationGenerator:

    def __init__(
        self,
        category: str = "finance",
        model: str = "gpt-4o-mini",
        temperature: float = 0.3,
        output_file: str = "financial_kb_output.json"
    ):
        self.category = category
        self.output_file = output_file

        self.llm = ChatOpenAI(
            model=model,
            temperature=temperature
        )

    # =========================
    # LOAD QUESTIONS FROM FILE
    # =========================
    def load_questions(self, file_path: str) -> List[str]:
        with open(file_path, "r") as f:
            lines = [line.strip() for line in f if line.strip()]

        print(f"✅ Loaded {len(lines)} questions from {file_path}")
        return lines

    # =========================
    # CLEAN QUESTION
    # =========================
    def clean_question(self, question: str) -> str:
        question = question.strip()

        # remove extra spaces before ?
        question = re.sub(r"\s+\?", "?", question)

        # ensure first letter capitalized
        question = question[0].upper() + question[1:]

        if not question.endswith("?"):
            question += "?"

        return question

    # =========================
    # KEYWORDS
    # =========================
    def extract_keywords(self, question: str) -> List[str]:
        words = re.sub(r"[^\w\s]", "", question.lower()).split()

        stop_words = {"what", "is", "the", "a", "an", "how", "does", "do"}
        keywords = [w for w in words if w not in stop_words]

        return keywords[:5]

    # =========================
    # VARIATIONS
    # =========================
    def generate_variations(self, question: str) -> List[str]:
        base = question.replace("?", "").lower()

        variations = [
            f"What is {base}?",
            f"Can you explain {base}?",
            f"How does {base} work?",
            f"What does {base} mean?",
            f"How do I understand {base}?",
            f"Can you give a simple explanation of {base}?"
        ]

        # Remove duplicates
        return list(set(variations))

    # =========================
    # LLM ANSWER
    # =========================
    def generate_answer(self, question: str) -> str:
        prompt = f"""
You are a financial expert.

Answer the question clearly and concisely.

Rules:
- Max 3 sentences
- Beginner friendly
- No fluff
- No repetition

Question:
{question}
"""

        response = self.llm.invoke(prompt)
        return response.content.strip()

    # =========================
    # BUILD RECORD
    # =========================
    def build_record(self, question: str, record_id: int) -> dict:
        clean_q = self.clean_question(question)

        return {
            "id": record_id,
            "category": self.category,
            "question": clean_q,
            "answer": self.generate_answer(clean_q),
            "keywords": self.extract_keywords(clean_q),
            "variations": self.generate_variations(clean_q)
        }

    # =========================
    # PROCESS ALL QUESTIONS
    # =========================
    def process_file(self, input_file: str):
        questions = self.load_questions("kb_unique_questions.txt")

        records = []

        for idx, q in enumerate(questions, start=1):
            try:
                record = self.build_record(q, idx)
                records.append(record)

                print(f"✅ Processed {idx}: {record['question']}")

                # Prevent rate limiting
                time.sleep(0.5)

            except Exception as e:
                print(f"❌ Failed at {idx}: {q} | Error: {e}")

        self.save_to_json(records)

    # =========================
    # SAVE ALL RECORDS
    # =========================
    def save_to_json(self, records: List[dict]):
        with open(self.output_file, "w") as f:
            json.dump(records, f, indent=2)
        print(f"💾 Saved {len(records)} records to {self.output_file}")


if __name__ == "__main__":
    load_dotenv()
    generator = VariationGenerator()
    generator.process_file("kb_unique_questions.txt")