playreel/backend/utils/common_llm.py

63 lines
2.2 KiB
Python

"""structured output 공용 래퍼.
OpenAI 호환 /chat/completions를 쓰는 프로바이더면 base_url만 바꿔 그대로 쓴다.
이미지는 이미 인코딩된 URL(data URI 또는 http)로 받는다 — 인코딩은 호출측 책임.
"""
from collections.abc import Sequence
from typing import TypeVar
from openai import AsyncOpenAI
from pydantic import BaseModel
T = TypeVar("T", bound=BaseModel)
class StructuredLLM:
def __init__(self, model: str, api_key: str, base_url: str | None = None):
self.model = model
self._client = AsyncOpenAI(api_key=api_key, base_url=base_url)
async def ask(
self,
schema: type[T],
prompt: str,
*,
system: str | None = None,
temperature: float = 0.0,
) -> T:
return await self._parse(schema, prompt, system, temperature)
async def ask_with_images(
self,
schema: type[T],
prompt: str,
images: Sequence[tuple[str, str]],
*,
detail: str = "high",
system: str | None = None,
temperature: float = 0.0,
) -> T:
"""images는 (라벨, 이미지 URL) 쌍. 라벨이 이미지 바로 앞에 붙어 순서가 보존된다."""
content: list[dict] = [{"type": "text", "text": prompt}]
for label, url in images:
content.append({"type": "text", "text": label})
content.append({"type": "image_url", "image_url": {"url": url, "detail": detail}})
return await self._parse(schema, content, system, temperature)
async def _parse(self, schema: type[T], content, system: str | None, temperature: float) -> T:
messages = [{"role": "system", "content": system}] if system else []
messages.append({"role": "user", "content": content})
msg = (await self._client.chat.completions.parse(
model=self.model,
messages=messages,
response_format=schema,
temperature=temperature,
)).choices[0].message
if msg.refusal:
raise RuntimeError(f"{self.model} 거절: {msg.refusal}")
if msg.parsed is None:
raise RuntimeError(f"{self.model} 파싱 실패: {msg.content!r}")
return msg.parsed