#!/usr/bin/env python3 """ Extract user segments from Reddit comments via OpenAI Responses API (Structured Outputs). Input: CSV file with columns: comment_id, post_url, comment Output: CSV file with the structure: comment_id, post_url, comment, segment, quotes Usage: export OPENAI_API_KEY=sk-... python extract_segments.py --input comments.csv --output segments.csv --model gpt-4o Notes: - Segments: Family, Freelancer, Student, Small Business, General, Unknown - Quotes are exact substrings from the comment that signal the segment - Uses batching & retries, just like pain_points/delighters """ import os import json import time import math import argparse import csv import hashlib import re from typing import List, Dict, Any def gen_id(text): return hashlib.sha256(text.encode("utf-8")).hexdigest() # OpenAI SDK try: from openai import OpenAI except Exception as e: raise SystemExit("Please install the official OpenAI Python SDK: pip install openai") from e # Optional tokenizer try: import tiktoken _ENC = tiktoken.get_encoding("cl100k_base") except Exception: _ENC = None # --------- Config Defaults --------- DEFAULT_MODEL = "gpt-4o" INPUT_TOKEN_BUDGET = 3000 OUTPUT_TOKENS = 800 TEMP = 0.1 MAX_RETRIES = 4 # --------- Schema --------- SEGMENT_SCHEMA = { "name": "segments", "schema": { "type":"object", "required":["items"], "properties":{ "items":{ "type":"array", "items":{ "type":"object", "required":["comment_id","segment"], "properties":{ "comment_id":{"type":"string"}, "segment":{"type":"string"}, # One of the categories "quotes":{"type":"array","items":{"type":"string"}} }, "additionalProperties":False } } }, "additionalProperties":False } } # --------- Prompts --------- SYSTEM_PROMPT = """You extract *user segments* from Reddit comments. Segments you can assign: - Family (mentions spouse, partner, kids, household, shared budgeting, family expenses) - Freelancer (mentions freelancing, gigs, invoices, clients, self-employed, independent work) - Student (mentions being a student, college, university, allowance, studying, campus life) - Small Business (mentions company, employees, team, payroll, business expenses, SME context) - General (default personal user, no special group cues) - Unknown (unclear / no information) Your tasks: - For each comment, decide which segment best fits. - Extract short verbatim substrings (quotes) that indicate the segment. Rules: - If multiple segments are hinted, choose the *most explicit* one. - If none apply, return "General" if it seems like a personal user, otherwise "Unknown". - Quotes must be verbatim substrings from the comment text (1–2 short phrases). """ USER_INSTRUCTIONS = """Process the following Reddit comments and return ONLY valid JSON matching the schema. For each comment: extract the most likely user segment and evidence quotes. Comments (array of objects): {comments_json} """ # --------- CSV Writer --------- def save_to_csv(all_items: list, comments_map: dict, output_path: str): fieldnames = ["comment_id", "post_url", "comment", "segment", "quotes"] with open(output_path, "w", encoding="utf-8", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() for item in all_items: if not isinstance(item, dict): continue cid = item.get("comment_id", "") original = comments_map.get(cid, {}) writer.writerow({ "comment_id": cid, "post_url": original.get("post_url", ""), "comment": original.get("text", ""), "segment": item.get("segment", ""), "quotes": " || ".join(item.get("quotes", [])) }) # --------- Utilities --------- def count_tokens(text: str) -> int: if _ENC is None: return max(1, math.ceil(len(text) / 4)) return len(_ENC.encode(text)) def pack_batches(comments: List[Dict[str, str]], max_total_tokens: int, expected_output_tokens_per_comment: int = 40) -> List[List[Dict[str, str]]]: batches = [] current = [] current_input_tokens = 0 overhead = 1000 for item in comments: s = json.dumps(item, ensure_ascii=False) t = count_tokens(s) est_output = (len(current) + 1) * expected_output_tokens_per_comment if t + current_input_tokens + overhead + est_output > max_total_tokens: if current: batches.append(current) current = [item] current_input_tokens = t else: current.append(item) current_input_tokens += t if current: batches.append(current) return batches def run_with_retries(fn, *args, **kwargs): delay = 2.0 for attempt in range(1, MAX_RETRIES + 1): try: return fn(*args, **kwargs) except Exception as e: if attempt == MAX_RETRIES: raise time.sleep(delay) delay = min(30.0, delay * 1.8) def safe_parse_response(resp): func_args = resp.choices[0].message.function_call.arguments if func_args: try: return json.loads(func_args) except Exception: m = re.search(r'\{.*\}|\[.*\]', func_args, re.DOTALL) if m: try: return json.loads(m.group(0)) except Exception: pass content = resp.choices[0].message.content if content: try: return json.loads(content) except Exception as e: raise RuntimeError(f"Could not parse model output:\n{content[:500]}...") from e raise RuntimeError("No valid JSON found in model response.") def call_openai_batch(client, model: str, batch: list) -> dict: user = USER_INSTRUCTIONS.format(comments_json=json.dumps(batch, ensure_ascii=False)) fc_schema = { "name": "extract_segments", "description": "Extract user segments from a batch of Reddit comments.", "parameters": SEGMENT_SCHEMA["schema"] } resp = client.chat.completions.create( model=model, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user} ], functions=[fc_schema], function_call={"name": "extract_segments"}, temperature=TEMP, max_tokens=OUTPUT_TOKENS ) return safe_parse_response(resp) def main(): ap = argparse.ArgumentParser() ap.add_argument("--input", required=True, help="Path to input csv with [{comment_id, text}, ...]") ap.add_argument("--output", required=True, help="Path to output csv with segments") ap.add_argument("--model", default=DEFAULT_MODEL, help=f"Model name (default: {DEFAULT_MODEL})") ap.add_argument("--batch_tokens", type=int, default=INPUT_TOKEN_BUDGET, help="Approx input token budget per request") args = ap.parse_args() api_key = os.getenv("OPENAI_API_KEY") if not api_key: raise SystemExit("Set OPENAI_API_KEY in your environment.") client = OpenAI(api_key=api_key) comments = [] not_processed_comments = [] with open(args.input, "r", encoding="utf-8") as f: reader = csv.DictReader(f) for row in reader: txt = str(row.get("comment", "")).strip() comment_id = gen_id(txt) url = str(row.get("post_url", "")).strip() if txt: comments.append({"comment_id": comment_id, "text": txt, "post_url": url}) batches = pack_batches(comments, args.batch_tokens) print(f"Processing {len(comments)} comments in {len(batches)} batches...") comments_map = {c["comment_id"]: c for c in comments} all_items = [] for i, batch in enumerate(batches, 1): print(f"Batch {i}/{len(batches)}: {len(batch)} comments") try: out = run_with_retries(call_openai_batch, client, args.model, batch) all_items.extend(out.get("items", [])) save_to_csv(all_items, comments_map, args.output) except RuntimeError: continue print(f"✅ All results saved to {args.output}") if __name__ == "__main__": main()