src/data_collection/collection.py

from dotenv import load_dotenv
import os
import pandas as pd
import time

import google.generativeai as genai

# ==========================================
# LOAD ENV VARIABLES
# ==========================================

load_dotenv()

# ==========================================
# CONFIGURE GEMINI
# ==========================================

genai.configure(
    api_key=os.getenv("GEMINI_API_KEY").strip()
)

# ==========================================
# LOAD PROMPT BANK
# ==========================================

df = pd.read_csv(
    "AI-MODEL-FINGERPRINTING/src/data_collection/prompt_bank.csv"
)

# ==========================================
# CLEAN COLUMN NAMES
# ==========================================

df.columns = df.columns.str.strip()

print("\nColumns:")
print(df.columns)

# ==========================================
# SELECT PROMPT RANGE
# ==========================================

# Example:
# 0:20   -> prompts 1 to 20
# 20:40  -> prompts 21 to 40
# 40:60  -> prompts 41 to 60

df = df.iloc[198:200]

# ==========================================
# OUTPUT FILE
# ==========================================



output_file = "AI-MODEL-FINGERPRINTING/src/data_collection/gemini_responses.csv"

# ==========================================
# LOAD EXISTING RESPONSES
# ==========================================

if os.path.exists(output_file):

    old_df = pd.read_csv(output_file)

    completed_ids = set(old_df["prompt_id"])

    print(f"\nAlready completed: {len(completed_ids)} prompts")

else:

    old_df = pd.DataFrame()

    completed_ids = set()

# ==========================================
# GEMINI RESPONSE FUNCTION
# ==========================================

def generate_gemini_response(prompt):

    try:

        model = genai.GenerativeModel(
            "models/gemini-2.5-flash"
        )

        response = model.generate_content(
            prompt
        )

        return response.text

    except Exception as e:

        print(f"\nGemini Error: {e}")

        return None

# ==========================================
# MAIN LOOP
# ==========================================

results = []

for index, row in df.iterrows():

    prompt_id = row["PROMPT_ID"]

    # ======================================
    # SKIP COMPLETED PROMPTS
    # ======================================

    if prompt_id in completed_ids:

        print(f"\nSkipping {prompt_id} (already completed)")

        continue

    category = row["category"]

    prompt = row["PROMPT"]

    # ======================================
    # STANDARDIZED PROMPT TEMPLATE
    # ======================================

    final_prompt = f"""
You are an advanced AI assistant.

Instructions:
- Respond ONLY in plain text.
- Do NOT use markdown.
- Do NOT use bullet points.
- Do NOT use numbered lists.
- Do NOT use headings.
- Avoid special formatting characters.
- Keep the tone natural and informative.
- Keep the response between 120 and 180 words.
- Give a complete and coherent answer.

USER PROMPT:
{prompt}
"""

    print(f"\nGenerating response for {prompt_id}...")

    # ======================================
    # GENERATE RESPONSE
    # ======================================

    generated_text = generate_gemini_response(
        final_prompt
    )

    # ======================================
    # HANDLE FAILED RESPONSES
    # ======================================

    if generated_text is None:

        print(f"\nFailed for {prompt_id}")

        continue

    # ======================================
    # STORE RESULT
    # ======================================

    result = {

        "prompt_id": prompt_id,

        "category": category,

        "model": "gemini",

        "prompt": prompt,

        "response": generated_text

    }

    results.append(result)

    # ======================================
    # SAVE AFTER EVERY RESPONSE
    # ======================================

    temp_df = pd.DataFrame(results)

    final_df = pd.concat(
        [old_df, temp_df],
        ignore_index=True
    )

    final_df.to_csv(
        output_file,
        index=False
    )

    print(f"\nSaved {prompt_id}")

    # ======================================
    # WAIT TO AVOID RATE LIMITS
    # ======================================

    time.sleep(3)

# ==========================================
# COMPLETION MESSAGE
# ==========================================

print("\nGemini response generation completed.")