""" S1 Re-index script: switch all tenant_internal thread embeddings to full-thread-text strategy. The A/B test (eval_embedding_ab.py) showed S1 (embed on full thread text = customer question + longest outbound reply) raises recall@3 from 5/13 → 7/13 vs. the short GPT summary used as the embedding input in production. model/dims unchanged (text-embedding-3-small @ 512), so no schema migration needed. This script re-embeds all existing rows using the new strategy WITHOUT re-generating the GPT summary (that's expensive and unnecessary — summary_text is for display, not the embedding input). Only the embedding vector changes. For threads without an existing summary, we fall back to the full-text strategy directly. Usage: venv/bin/python scripts/reindex_s1.py [--schema tenant_internal] [--dry-run] """ from __future__ import annotations import argparse import sys import time from pathlib import Path from typing import Dict, List, Optional, Tuple sys.path.insert(0, str(Path(__file__).resolve().parents[1])) USD_EUR = 0.92 PRICE_USD_PER_1M_EMBED = 0.02 # text-embedding-3-small def run_reindex(schema: str = "tenant_internal", dry_run: bool = False) -> None: print(f"\n{'='*60}") print(f"S1 Re-index: switch to full thread text embeddings") print(f"Schema: {schema} | dry_run={dry_run}") print(f"{'='*60}\n") from supporthub.app.config import config from supporthub.app.db import session_scope, set_worker_tenant_schema from supporthub.app.services.embedding_service import ( EmbeddingService, pack_embedding, compute_text_hash, ) from sqlalchemy import text import struct if not config.AI_API_KEY: print("ERROR: AI_API_KEY not set.") sys.exit(1) set_worker_tenant_schema(schema) svc = EmbeddingService() # ----------------------------------------------------------------------- # 1. Load all thread_embeddings rows that have an existing blob # ----------------------------------------------------------------------- with session_scope(schema=schema) as session: rows = session.execute(text(""" SELECT thread_id, summary_text, dimensions, product_id, embedding_quality FROM thread_embeddings WHERE embedding_blob IS NOT NULL ORDER BY thread_id """)).fetchall() print(f"Existing rows to re-embed: {len(rows)}") # ----------------------------------------------------------------------- # 2. For each thread, load first inbound + longest outbound from DB # ----------------------------------------------------------------------- print("Loading inbound + longest outbound for each thread...") thread_ids = [r.thread_id for r in rows] row_meta: Dict[int, Dict] = { r.thread_id: { "summary_text": r.summary_text or "", "dimensions": r.dimensions, "product_id": r.product_id, "quality": r.embedding_quality, } for r in rows } # Batch load first inbound per thread CHUNK = 1000 inbound_map: Dict[int, str] = {} longest_outbound_map: Dict[int, str] = {} for i in range(0, len(thread_ids), CHUNK): chunk = thread_ids[i:i + CHUNK] with session_scope(schema=schema) as session: inbound_rows = session.execute(text(""" SELECT DISTINCT ON (thread_id) thread_id, body_text FROM addons_messages WHERE thread_id = ANY(:ids) AND direction = 'inbound' ORDER BY thread_id, created_at ASC """), {"ids": chunk}).fetchall() for r in inbound_rows: if r.body_text: inbound_map[r.thread_id] = r.body_text longest_rows = session.execute(text(""" SELECT DISTINCT ON (thread_id) thread_id, body_text FROM addons_messages WHERE thread_id = ANY(:ids) AND direction = 'outbound' AND body_text IS NOT NULL ORDER BY thread_id, length(body_text) DESC """), {"ids": chunk}).fetchall() for r in longest_rows: if r.body_text: longest_outbound_map[r.thread_id] = r.body_text if (i // CHUNK + 1) % 5 == 0: print(f" Loaded chunk {i}-{i+CHUNK}...") print(f" Inbound available for {len(inbound_map)} threads") print(f" Outbound available for {len(longest_outbound_map)} threads\n") # ----------------------------------------------------------------------- # 3. Build full thread texts for embedding # ----------------------------------------------------------------------- to_embed: List[Tuple[int, str]] = [] # (thread_id, full_text) skipped_no_content = 0 for tid in thread_ids: inbound = inbound_map.get(tid, "") longest_out = longest_outbound_map.get(tid, "") if not inbound: skipped_no_content += 1 continue full_text = svc._build_full_thread_text(inbound, longest_out) to_embed.append((tid, full_text)) print(f"Threads to embed: {len(to_embed)} (skipped {skipped_no_content} with no inbound)") # Estimate cost total_chars = sum(len(t) for _, t in to_embed) est_tokens = total_chars // 4 # rough chars-per-token estimate est_usd = est_tokens / 1_000_000 * PRICE_USD_PER_1M_EMBED est_eur = est_usd * USD_EUR print(f"Estimated embedding cost: ~{est_tokens/1_000_000:.2f}M tokens = ~€{est_eur:.3f}\n") if dry_run: print("[DRY RUN] Would re-embed these threads. No changes made.") return # ----------------------------------------------------------------------- # 4. Batch embed and update DB # ----------------------------------------------------------------------- import openai as _openai BATCH_SIZE = 100 total_embedded = 0 total_tokens = 0 t0 = time.monotonic() for i in range(0, len(to_embed), BATCH_SIZE): batch = to_embed[i:i + BATCH_SIZE] batch_ids = [tid for tid, _ in batch] batch_texts = [txt[:8000] for _, txt in batch] try: resp = svc.client.embeddings.create( input=batch_texts, model=svc.model, dimensions=svc.dimensions, ) except Exception as exc: print(f" ERROR batch {i}-{i+BATCH_SIZE}: {exc}") continue if resp.usage: total_tokens += resp.usage.total_tokens sorted_data = sorted(resp.data, key=lambda d: d.index) # Store results with session_scope(schema=schema) as session: for j, emb_data in enumerate(sorted_data): tid = batch_ids[j] vec = emb_data.embedding blob = pack_embedding(vec) meta = row_meta.get(tid, {}) session.execute(text(""" UPDATE thread_embeddings SET embedding_blob = :blob, model_used = :model, dimensions = :dims, text_hash = :hash WHERE thread_id = :tid """), { "blob": blob, "model": svc.model, "dims": svc.dimensions, "hash": compute_text_hash(batch_texts[j]), "tid": tid, }) total_embedded += 1 elapsed = time.monotonic() - t0 pct = (i + len(batch)) / len(to_embed) * 100 cost_so_far = total_tokens / 1_000_000 * PRICE_USD_PER_1M_EMBED * USD_EUR print(f" [{pct:.0f}%] {total_embedded}/{len(to_embed)} embedded | " f"€{cost_so_far:.4f} | {elapsed:.0f}s elapsed") elapsed_total = time.monotonic() - t0 actual_cost = total_tokens / 1_000_000 * PRICE_USD_PER_1M_EMBED * USD_EUR print(f"\n{'='*60}") print(f"S1 Re-index complete") print(f" Threads re-embedded: {total_embedded}") print(f" Total tokens: {total_tokens:,}") print(f" Actual cost: €{actual_cost:.4f}") print(f" Time: {elapsed_total:.1f}s") print(f"{'='*60}\n") if __name__ == "__main__": parser = argparse.ArgumentParser(description="S1 re-index: switch to full thread text embeddings") parser.add_argument("--schema", default="tenant_internal", help="Tenant schema") parser.add_argument("--dry-run", action="store_true", help="Estimate cost without making changes") args = parser.parse_args() run_reindex(schema=args.schema, dry_run=args.dry_run)