o2o-castad-backend/app/utils/prompts/prompts.py

124 lines
4.2 KiB
Python

import re
import gspread
from pydantic import BaseModel
from google.oauth2.service_account import Credentials
from config import prompt_settings
from app.utils.logger import get_logger
from app.utils.prompts.schemas import *
from functools import lru_cache
# {known_field} 형태만 치환하기 위한 패턴. JSON 예시 등 리터럴 중괄호는 보존된다.
_PLACEHOLDER_RE = re.compile(r"\{(\w+)\}")
logger = get_logger("prompt")
_SCOPES = [
"https://www.googleapis.com/auth/spreadsheets.readonly"
]
_sheet_cache: dict[str, tuple[str, str]] = {}
def _get_spreadsheet():
creds = Credentials.from_service_account_file(
prompt_settings.GOOGLE_SERVICE_ACCOUNT_JSON, scopes=_SCOPES
)
return gspread.authorize(creds).open_by_key(prompt_settings.PROMPT_SPREADSHEET)
def _preload_all_sheets(sheet_names: list[str]):
ranges = [f"{name}!B2:B3" for name in sheet_names]
response = _get_spreadsheet().values_batch_get(ranges)
for value_range in response["valueRanges"]:
range_name = value_range.get("range", "")
name = range_name.split("!")[0].strip("'")
if name not in sheet_names:
continue
values = value_range.get("values", [])
if len(values) < 2 or not values[0] or not values[1]:
logger.warning(f"Sheet '{name}' has missing data, skipping preload")
continue
_sheet_cache[name] = (values[1][0], values[0][0]) # (B3=template, B2=model)
def _read_sheet_data(sheet_name: str) -> tuple[str, str]:
if sheet_name not in _sheet_cache:
ws = _get_spreadsheet().worksheet(sheet_name)
values = ws.batch_get(["B2:B3"])[0]
_sheet_cache[sheet_name] = (values[1][0], values[0][0])
return _sheet_cache[sheet_name]
class Prompt():
sheet_name: str
prompt_template: str
prompt_model: str
prompt_input_class = BaseModel
prompt_output_class = BaseModel
def __init__(self, sheet_name, prompt_input_class, prompt_output_class):
self.sheet_name = sheet_name
self.prompt_input_class = prompt_input_class
self.prompt_output_class = prompt_output_class
self.prompt_template, self.prompt_model = _read_sheet_data(sheet_name)
def build_prompt(self, input_data:dict, silent:bool = False) -> str:
verified_input = self.prompt_input_class(**input_data)
values = verified_input.model_dump()
# 알려진 {필드}만 치환하고, 프롬프트 내 JSON 예시 등 리터럴 중괄호는 그대로 둔다.
# (str.format은 모든 {..}를 필드로 해석해 리터럴 중괄호에서 깨지므로 사용하지 않음)
build_template = _PLACEHOLDER_RE.sub(
lambda m: str(values[m.group(1)]) if m.group(1) in values else m.group(0),
self.prompt_template,
)
if not silent:
logger.debug(f"build_template: {build_template}")
logger.debug(f"input_data: {input_data}")
return build_template
# 업종 분기는 각 프롬프트 내부에서 {industry} 변수로 처리하므로 시트는 타입별 1개로 통합.
_preload_all_sheets([
"marketing",
"lyric",
"subtitle",
"yt_upload",
"image_tag",
])
# 업종 분기는 프롬프트 내부 {industry} 변수로 처리하므로 타입별 단일 프롬프트만 둔다.
marketing_prompt = Prompt(
sheet_name="marketing",
prompt_input_class=MarketingPromptInput,
prompt_output_class=MarketingPromptOutput,
)
lyric_prompt = Prompt(
sheet_name="lyric",
prompt_input_class=LyricPromptInput,
prompt_output_class=LyricPromptOutput,
)
yt_upload_prompt = Prompt(
sheet_name="yt_upload",
prompt_input_class=YTUploadPromptInput,
prompt_output_class=YTUploadPromptOutput,
)
image_autotag_prompt = Prompt(
sheet_name="image_tag",
prompt_input_class=ImageTagPromptInput,
prompt_output_class=ImageTagPromptOutput,
)
@lru_cache()
def create_dynamic_subtitle_prompt(length: int, industry: str = "") -> Prompt:
# industry 인자는 캐시 구분/하위 호환용. 시트는 단일 'subtitle'로 통합됨.
return Prompt(
sheet_name="subtitle",
prompt_input_class=SubtitlePromptInput,
prompt_output_class=SubtitlePromptOutput[length],
)