data_gen.py

from openai import OpenAI
from pydantic import BaseModel
from typing import List
import json
import os
from dotenv import load_dotenv

load_dotenv()  # Load environment variables from .env file

openai_key=os.getenv("OPENAI_API_KEY")
if not openai_key:
    raise ValueError("OPENAI_API_KEY not found in environment variables.")


# -------------------------------
# Initialize OpenAI client
# -------------------------------
# Make sure OPENAI_API_KEY is set in your environment before running
client=OpenAI(api_key=openai_key)
# -------------------------------
# Define Structured Output Schema
# -------------------------------
class Triplet(BaseModel):
    base: str
    looks_similar: str
    looks_different: str
    explanation: str

class TripletBatch(BaseModel):
    domain: str
    triplets: List[Triplet]

# -------------------------------
# Load Prompt from Markdown File
# -------------------------------
def load_prompt(filepath: str) -> str:
    with open(filepath, "r", encoding="utf-8") as f:
        return f.read()

# -------------------------------
# Generate Triplets (Single Batch)
# -------------------------------
def generate_triplets(prompt_text: str) -> TripletBatch:
    response = client.responses.parse(
        model="gpt-4.1",
        input=[
            {
                "role": "system",
                "content": (
                    "You are a linguistic data generator. "
                    "Produce high-quality semantic stress-test triplets "
                    "that distinguish word-level from meaning-level embeddings."
                ),
            },
            {"role": "user", "content": prompt_text},
        ],
        text_format=TripletBatch,
    )
    return response.output_parsed

# -------------------------------
# Generate Multiple Batches & Stream Save
# -------------------------------
def generate_and_save(prompt_path: str, n_batches: int, output_file: str):
    prompt = load_prompt(prompt_path)

    # Open file in append mode so data is written immediately
    with open(output_file, "a", encoding="utf-8") as f:
        for i in range(1, n_batches + 1):
            print(f"🔹 Generating batch {i}/{n_batches}...")

            try:
                batch = generate_triplets(prompt)
            except Exception as e:
                print(f"⚠️ Error on batch {i}: {e}")
                continue  # Skip failed batch and move on

            for t in batch.triplets:
                triplet_dict = {
                    "Base": t.base,
                    "Looks Similar": t.looks_similar,
                    "Looks Different": t.looks_different,
                    "Explanation": t.explanation,
                }
                json.dump(triplet_dict, f, ensure_ascii=False)
                f.write("\n")
                f.flush()  # Write immediately, don’t buffer

            print(f"✅ Saved batch {i} to file (streamed).")

    print(f"\n🎯 All done! {n_batches} batches written to '{output_file}'")

# -------------------------------
# Main
# ------------------------------
if __name__ == "__main__":
    generate_and_save(
        prompt_path="software_repo_desc/software_gen_triplet.md",
     n_batches=100,
      output_file="software_repo_desc/triplets.jsonl")