""" Spot Embedding Generator Gemini text-embedding-004 모델을 사용하여 story_spots의 벡터 임베딩을 생성하고 Supabase spot_embeddings 테이블에 저장합니다. Usage: from utils.embedding_generator import generate_all_embeddings, update_changed_embeddings # 전체 생성 (초기 1회) stats = await generate_all_embeddings() # 변경분만 업데이트 (주기적) stats = await update_changed_embeddings() # CLI 직접 실행 python -m utils.embedding_generator [--force] """ import asyncio import hashlib import json import logging import os import time from typing import Any, Optional from google import genai from google.genai import types logger = logging.getLogger(__name__) # ============ Constants ============ EMBEDDING_MODEL = "text-embedding-004" EMBEDDING_DIMENSIONS = 768 BATCH_SIZE = 100 # Gemini batch embed limit RATE_LIMIT_DELAY = 0.5 # seconds between batch calls # Gemini task types for embedding TASK_TYPE_RETRIEVAL_DOCUMENT = "RETRIEVAL_DOCUMENT" TASK_TYPE_RETRIEVAL_QUERY = "RETRIEVAL_QUERY" # story_spots columns needed for embedding text construction SPOT_SELECT_COLUMNS = ( "id, name, name_en, name_zh, category, lat, lng, address, " "story_title, story_content, story_source, tips, " "tags_tier1, tags_tier2, meta, " "main_image_url, thumbnail_url, generated_image_url, " "priority_score, status, zone, cluster_id, " "village, source_book, historical_period" ) # ============ Gemini Client ============ def _get_gemini_client() -> genai.Client: """Get Gemini API client (lazy init).""" api_key = os.getenv("GEMINI_API_KEY") if not api_key: raise ValueError("GEMINI_API_KEY environment variable is required") return genai.Client(api_key=api_key) # ============ Content Hash ============ def compute_content_hash(spot: dict[str, Any]) -> str: """ 스팟 데이터의 콘텐츠 해시 생성. 임베딩에 사용되는 필드만 포함하여, 관련 없는 필드 변경으로 불필요한 재생성 방지. """ hash_fields = { "name": spot.get("name", ""), "category": spot.get("category", ""), "story_title": spot.get("story_title", ""), "story_content": spot.get("story_content", ""), "tips": spot.get("tips", ""), "tags_tier1": json.dumps(spot.get("tags_tier1") or {}, sort_keys=True, ensure_ascii=False), "tags_tier2": json.dumps(spot.get("tags_tier2") or [], sort_keys=True, ensure_ascii=False), "village": spot.get("village", ""), "historical_period": spot.get("historical_period", ""), } content_str = json.dumps(hash_fields, sort_keys=True, ensure_ascii=False) return hashlib.sha256(content_str.encode("utf-8")).hexdigest()[:16] # ============ Embedding Text Construction ============ def build_embedding_text(spot: dict[str, Any]) -> str: """ 스팟 데이터를 임베딩용 텍스트로 변환. 구성: - 이름 (한국어 + 영어/중국어) - 카테고리 - 스토리 제목 + 내용 (핵심 시맨틱 정보) - 태그 (테마, 분위기, 활동성) - 팁 - 마을, 시대 정보 """ parts: list[str] = [] # 이름 (다국어) name = spot.get("name", "") parts.append(f"장소: {name}") if spot.get("name_en"): parts.append(f"({spot['name_en']})") if spot.get("name_zh"): parts.append(f"({spot['name_zh']})") # 카테고리 category = spot.get("category", "") if category: parts.append(f"분류: {category}") # 마을/시대 (향토지 메타데이터) if spot.get("village"): parts.append(f"마을: {spot['village']}") if spot.get("historical_period"): parts.append(f"시대: {spot['historical_period']}") # 스토리 (핵심 시맨틱) story_title = spot.get("story_title", "") story_content = spot.get("story_content", "") if story_title: parts.append(f"이야기: {story_title}") if story_content and "카테고리에 속합니다" not in story_content: # 자동 생성된 의미없는 콘텐츠 제외 # 임베딩 토큰 예산: ~500자로 제한 content_trimmed = story_content[:500] parts.append(content_trimmed) # 태그 tags_tier1 = spot.get("tags_tier1") or {} themes = tags_tier1.get("theme", []) moods = tags_tier1.get("mood", []) activity = tags_tier1.get("activity_level", "") time_of_day = tags_tier1.get("time_of_day", []) tag_parts: list[str] = [] if themes: tag_parts.append(f"테마: {', '.join(themes)}") if moods: tag_parts.append(f"분위기: {', '.join(moods)}") if activity: tag_parts.append(f"활동: {activity}") if time_of_day: tag_parts.append(f"시간대: {', '.join(time_of_day)}") if tag_parts: parts.append(" | ".join(tag_parts)) # tier2 태그 tags_tier2 = spot.get("tags_tier2") or [] if tags_tier2: parts.append(f"키워드: {', '.join(tags_tier2)}") # 팁 tips = spot.get("tips", "") if tips: parts.append(f"팁: {tips[:200]}") return "\n".join(parts) # ============ Embedding API ============ async def generate_embeddings_batch( texts: list[str], client: Optional[genai.Client] = None, task_type: str = TASK_TYPE_RETRIEVAL_DOCUMENT, ) -> list[list[float]]: """ Gemini text-embedding-004로 텍스트 배치 임베딩 생성. Args: texts: 임베딩할 텍스트 리스트 (최대 BATCH_SIZE) client: Gemini 클라이언트 (None이면 자동 생성) task_type: RETRIEVAL_DOCUMENT (저장용) 또는 RETRIEVAL_QUERY (검색용) Returns: 임베딩 벡터 리스트 (각 768차원) """ if not texts: return [] if client is None: client = _get_gemini_client() if len(texts) > BATCH_SIZE: raise ValueError(f"Batch size {len(texts)} exceeds limit {BATCH_SIZE}") try: response = await asyncio.to_thread( client.models.embed_content, model=EMBEDDING_MODEL, contents=texts, config=types.EmbedContentConfig( task_type=task_type, output_dimensionality=EMBEDDING_DIMENSIONS, ), ) embeddings = [] for emb in response.embeddings: embeddings.append(list(emb.values)) logger.info( f"[embedding] Generated {len(embeddings)} embeddings " f"(dim={len(embeddings[0]) if embeddings else 0})" ) return embeddings except Exception as e: logger.error(f"[embedding] Gemini embed_content failed: {e}") raise async def generate_query_embedding( query_text: str, client: Optional[genai.Client] = None, ) -> list[float]: """ 검색 쿼리를 임베딩으로 변환. RETRIEVAL_QUERY task type 사용 (문서 임베딩과 쌍을 이룸). Args: query_text: 검색 쿼리 텍스트 Returns: 768차원 임베딩 벡터 """ embeddings = await generate_embeddings_batch( [query_text], client=client, task_type=TASK_TYPE_RETRIEVAL_QUERY, ) return embeddings[0] # ============ DB Operations ============ def _load_active_spots() -> list[dict[str, Any]]: """Supabase에서 active 스팟 로드 (임베딩 생성에 필요한 컬럼만).""" from db import get_supabase supabase = get_supabase() result = supabase.table("story_spots") \ .select(SPOT_SELECT_COLUMNS) \ .eq("status", "active") \ .execute() return result.data or [] def _load_existing_embeddings() -> dict[str, str]: """기존 임베딩의 spot_id -> content_hash 매핑 로드.""" from db import get_supabase supabase = get_supabase() result = supabase.table("spot_embeddings") \ .select("spot_id, content_hash") \ .execute() return {row["spot_id"]: row["content_hash"] for row in (result.data or [])} def _upsert_embeddings(records: list[dict[str, Any]]) -> int: """임베딩 레코드 배치 upsert. 성공 건수 반환.""" if not records: return 0 from db import get_supabase supabase = get_supabase() # Supabase upsert: spot_id UNIQUE constraint로 ON CONFLICT result = supabase.table("spot_embeddings") \ .upsert(records, on_conflict="spot_id") \ .execute() return len(result.data) if result.data else 0 def _delete_orphan_embeddings(active_spot_ids: set[str]) -> int: """더 이상 active가 아닌 스팟의 임베딩 삭제.""" from db import get_supabase supabase = get_supabase() # 현재 임베딩이 있는 spot_id 조회 result = supabase.table("spot_embeddings") \ .select("spot_id") \ .execute() orphan_ids = [ row["spot_id"] for row in (result.data or []) if row["spot_id"] not in active_spot_ids ] if not orphan_ids: return 0 for orphan_id in orphan_ids: supabase.table("spot_embeddings") \ .delete() \ .eq("spot_id", orphan_id) \ .execute() logger.info(f"[embedding] Deleted {len(orphan_ids)} orphan embeddings") return len(orphan_ids) # ============ Main Functions ============ async def generate_all_embeddings(force: bool = False) -> dict[str, Any]: """ 전체 스팟 임베딩 생성. Args: force: True면 content_hash 무시하고 전부 재생성 Returns: { "total_spots": int, "generated": int, "skipped": int, "deleted_orphans": int, "errors": int, "duration_seconds": float } """ start = time.monotonic() stats = { "total_spots": 0, "generated": 0, "skipped": 0, "deleted_orphans": 0, "errors": 0, "duration_seconds": 0.0, } try: # 1. Active 스팟 로드 spots = await asyncio.to_thread(_load_active_spots) stats["total_spots"] = len(spots) logger.info(f"[embedding] Loaded {len(spots)} active spots") if not spots: return stats # 2. 기존 임베딩 해시 로드 existing_hashes = await asyncio.to_thread(_load_existing_embeddings) logger.info(f"[embedding] Found {len(existing_hashes)} existing embeddings") # 3. 변경 감지: 새로운/변경된 스팟만 필터링 spots_to_embed: list[dict[str, Any]] = [] for spot in spots: content_hash = compute_content_hash(spot) spot["_content_hash"] = content_hash if force or spot["id"] not in existing_hashes: spots_to_embed.append(spot) elif existing_hashes[spot["id"]] != content_hash: spots_to_embed.append(spot) else: stats["skipped"] += 1 logger.info( f"[embedding] To embed: {len(spots_to_embed)}, " f"skipped (unchanged): {stats['skipped']}" ) # 4. 배치 임베딩 생성 client = _get_gemini_client() for batch_start in range(0, len(spots_to_embed), BATCH_SIZE): batch = spots_to_embed[batch_start:batch_start + BATCH_SIZE] batch_texts = [build_embedding_text(s) for s in batch] try: embeddings = await generate_embeddings_batch( batch_texts, client=client ) # DB upsert 레코드 준비 records = [] for spot, embedding in zip(batch, embeddings): records.append({ "spot_id": spot["id"], "embedding": embedding, "embedding_model": EMBEDDING_MODEL, "content_hash": spot["_content_hash"], "updated_at": "now()", }) upserted = await asyncio.to_thread(_upsert_embeddings, records) stats["generated"] += upserted logger.info( f"[embedding] Batch {batch_start // BATCH_SIZE + 1}: " f"embedded {upserted}/{len(batch)} spots" ) except Exception as e: stats["errors"] += len(batch) logger.error( f"[embedding] Batch {batch_start // BATCH_SIZE + 1} failed: {e}" ) # Rate limiting if batch_start + BATCH_SIZE < len(spots_to_embed): await asyncio.sleep(RATE_LIMIT_DELAY) # 5. 고아 임베딩 정리 active_ids = {s["id"] for s in spots} stats["deleted_orphans"] = await asyncio.to_thread( _delete_orphan_embeddings, active_ids ) except Exception as e: logger.error(f"[embedding] generate_all_embeddings failed: {e}") stats["errors"] += 1 stats["duration_seconds"] = round(time.monotonic() - start, 2) logger.info(f"[embedding] Complete: {stats}") return stats async def update_changed_embeddings() -> dict[str, Any]: """ 변경된 스팟만 임베딩 업데이트. content_hash 비교로 실제 변경분만 처리. generate_all_embeddings(force=False)의 편의 래퍼. """ return await generate_all_embeddings(force=False) # ============ CLI Entry Point ============ async def _main(): """CLI 실행용 메인 함수.""" import argparse import sys # backend/ 디렉토리를 Python path에 추가 from pathlib import Path backend_dir = str(Path(__file__).parent.parent) if backend_dir not in sys.path: sys.path.insert(0, backend_dir) from dotenv import load_dotenv load_dotenv() logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", ) parser = argparse.ArgumentParser(description="Generate spot embeddings") parser.add_argument( "--force", action="store_true", help="Force regenerate all embeddings (ignore content_hash)" ) parser.add_argument( "--dry-run", action="store_true", help="Show what would be embedded without actually calling API" ) args = parser.parse_args() if args.dry_run: spots = _load_active_spots() existing = _load_existing_embeddings() changed = 0 new = 0 for spot in spots: h = compute_content_hash(spot) if spot["id"] not in existing: new += 1 print(f" NEW: {spot['id']} ({spot['name']})") elif existing[spot["id"]] != h: changed += 1 print(f" CHANGED: {spot['id']} ({spot['name']})") print(f"\nTotal: {len(spots)} active spots") print(f" New: {new}") print(f" Changed: {changed}") print(f" Unchanged: {len(spots) - new - changed}") return stats = await generate_all_embeddings(force=args.force) print(f"\nEmbedding generation complete:") for k, v in stats.items(): print(f" {k}: {v}") if __name__ == "__main__": asyncio.run(_main())