Initial commit
This commit is contained in:
jsahn
@@ -0,0 +1,305 @@
|
||||
"""MiniMax provider: OpenAI-compatible chat completions with JSON-schema output and model fallback."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Any, get_origin
|
||||
|
||||
import openai
|
||||
from pydantic import ValidationError
|
||||
|
||||
from us2kr.config import Settings
|
||||
from us2kr.llm.base import LLMError
|
||||
from us2kr.llm.prompts import SYSTEM_PROMPT, build_user_message
|
||||
from us2kr.models import MarketPayload, Report
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Errors from the openai SDK that mean "this model/attempt failed; try the next model".
|
||||
OPENAI_ERRORS: tuple[type[Exception], ...] = (
|
||||
openai.APIStatusError,
|
||||
openai.APIConnectionError,
|
||||
openai.APITimeoutError,
|
||||
openai.RateLimitError,
|
||||
)
|
||||
|
||||
_THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL)
|
||||
_CODE_FENCE = re.compile(r"```(?:json|JSON)?\s*(.*?)```", re.DOTALL)
|
||||
|
||||
REPAIR_INSTRUCTION = (
|
||||
"위 응답은 JSON 스키마 검증에 실패했습니다. 아래 오류를 참고하여, 설명 없이 스키마에 맞는 "
|
||||
"완전한 JSON 객체 하나만 다시 출력하세요.\n\n오류:\n"
|
||||
)
|
||||
|
||||
|
||||
def _strip_titles(node: Any) -> None:
|
||||
"""Recursively remove JSON-schema "title" keys (noise for the model, not needed for validation)."""
|
||||
if isinstance(node, dict):
|
||||
node.pop("title", None)
|
||||
for value in node.values():
|
||||
_strip_titles(value)
|
||||
elif isinstance(node, list):
|
||||
for value in node:
|
||||
_strip_titles(value)
|
||||
|
||||
|
||||
def report_json_schema() -> dict[str, Any]:
|
||||
"""JSON schema for `Report` suitable for `response_format` (no titles, closed object)."""
|
||||
schema = copy.deepcopy(Report.model_json_schema())
|
||||
_strip_titles(schema)
|
||||
schema["additionalProperties"] = False
|
||||
return schema
|
||||
|
||||
|
||||
def report_response_format() -> dict[str, Any]:
|
||||
return {"type": "json_schema", "json_schema": {"name": "report", "schema": report_json_schema()}}
|
||||
|
||||
|
||||
def _first_json_object(text: str) -> str | None:
|
||||
"""Return the first balanced `{...}` block in `text` (string-aware), or None."""
|
||||
start = text.find("{")
|
||||
if start < 0:
|
||||
return None
|
||||
depth = 0
|
||||
in_str = False
|
||||
escape = False
|
||||
for i in range(start, len(text)):
|
||||
ch = text[i]
|
||||
if in_str:
|
||||
if escape:
|
||||
escape = False
|
||||
elif ch == "\\":
|
||||
escape = True
|
||||
elif ch == '"':
|
||||
in_str = False
|
||||
continue
|
||||
if ch == '"':
|
||||
in_str = True
|
||||
elif ch == "{":
|
||||
depth += 1
|
||||
elif ch == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return text[start : i + 1]
|
||||
return None
|
||||
|
||||
|
||||
def extract_json_text(raw: str) -> str:
|
||||
"""Strip <think> blocks and code fences from a model reply, returning the JSON candidate text.
|
||||
|
||||
Handles: complete <think>...</think> blocks anywhere, an unterminated <think> (take text after the
|
||||
last </think>, or the first {...} object), and ```json fences. Raises ValueError if nothing remains.
|
||||
"""
|
||||
text = _THINK_BLOCK.sub("", raw)
|
||||
if "<think>" in text:
|
||||
if "</think>" in raw:
|
||||
text = raw.rsplit("</think>", 1)[1]
|
||||
else:
|
||||
text = _first_json_object(raw) or ""
|
||||
fence = _CODE_FENCE.search(text)
|
||||
if fence:
|
||||
text = fence.group(1)
|
||||
text = _first_json_object(text.strip()) or ""
|
||||
if not text:
|
||||
raise ValueError("no JSON object found in model reply")
|
||||
return text
|
||||
|
||||
|
||||
_LIST_FIELDS: frozenset[str] = frozenset(
|
||||
name for name, info in Report.model_fields.items() if get_origin(info.annotation) is list
|
||||
)
|
||||
_BULLET_SPLIT = re.compile(r"(?:^|\n)\s*(?:[-*•]|\d+[.)])\s+")
|
||||
_SENTENCE_SPLIT = re.compile(r"(?<=[.!?。])\s+")
|
||||
|
||||
|
||||
def _split_prose(text: str) -> list[str]:
|
||||
"""Split a prose string into list items: bullets/newlines first, then sentence boundaries."""
|
||||
parts = [p.strip() for p in _BULLET_SPLIT.split(text) if p.strip()]
|
||||
if len(parts) <= 1:
|
||||
parts = [p.strip() for p in re.split(r"\n+", text) if p.strip()]
|
||||
if len(parts) <= 1:
|
||||
parts = [p.strip() for p in _SENTENCE_SPLIT.split(text) if p.strip()]
|
||||
return parts or [text.strip()]
|
||||
|
||||
|
||||
def coerce_report_data(data: Any) -> tuple[Any, list[str]]:
|
||||
"""Fix the common MiniMax schema slip of emitting list fields as a single string.
|
||||
|
||||
Returns (data, names of coerced fields). Anything that is not a dict is returned untouched so that
|
||||
pydantic can report the real error.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
return data, []
|
||||
coerced: list[str] = []
|
||||
for name in _LIST_FIELDS:
|
||||
value = data.get(name)
|
||||
if isinstance(value, str):
|
||||
data[name] = _split_prose(value)
|
||||
coerced.append(name)
|
||||
return data, coerced
|
||||
|
||||
|
||||
def parse_report(raw: str) -> Report:
|
||||
"""Turn raw model content into a validated Report (raises ValueError / ValidationError)."""
|
||||
text = extract_json_text(raw)
|
||||
try:
|
||||
data = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
return Report.model_validate_json(text) # let pydantic produce the canonical error
|
||||
data, coerced = coerce_report_data(data)
|
||||
if coerced:
|
||||
log.warning("coerced string -> list for report fields: %s", ", ".join(sorted(coerced)))
|
||||
return Report.model_validate(data)
|
||||
|
||||
|
||||
class MiniMaxProvider:
|
||||
"""LLMProvider backed by MiniMax's OpenAI-compatible endpoint with a model fallback chain."""
|
||||
|
||||
def __init__(self, settings: Settings, client: Any | None = None) -> None:
|
||||
self._settings = settings
|
||||
models = [settings.minimax_model, *settings.minimax_fallback_models]
|
||||
self._models = list(dict.fromkeys(m for m in models if m))
|
||||
if not self._models:
|
||||
raise LLMError("no MiniMax model configured")
|
||||
if client is None:
|
||||
if not settings.minimax_api_key:
|
||||
raise LLMError("MINIMAX_API_KEY is not set")
|
||||
client = openai.OpenAI(
|
||||
api_key=settings.minimax_api_key,
|
||||
base_url=settings.minimax_base_url,
|
||||
timeout=settings.minimax_timeout_seconds,
|
||||
max_retries=2,
|
||||
)
|
||||
self._client = client
|
||||
self._response_format = report_response_format()
|
||||
|
||||
@property
|
||||
def models(self) -> list[str]:
|
||||
return list(self._models)
|
||||
|
||||
# ------------------------------------------------------------------ public
|
||||
def summarize(self, payload: MarketPayload) -> tuple[Report, str]:
|
||||
user_message = build_user_message(payload)
|
||||
base_messages: list[dict[str, Any]] = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": user_message},
|
||||
]
|
||||
attempts: list[str] = []
|
||||
for model in self._models:
|
||||
result = self._attempt_model(model, base_messages, attempts)
|
||||
if result is not None:
|
||||
return result, model
|
||||
raise LLMError("all MiniMax models failed: " + " | ".join(attempts))
|
||||
|
||||
# ----------------------------------------------------------------- helpers
|
||||
def _attempt_model(
|
||||
self, model: str, messages: list[dict[str, Any]], attempts: list[str]
|
||||
) -> Report | None:
|
||||
"""One model: initial call plus at most one repair retry. Returns None when the model failed."""
|
||||
content = self._call(model, messages, attempts)
|
||||
if content is None:
|
||||
return None
|
||||
try:
|
||||
return parse_report(content)
|
||||
except (ValidationError, ValueError) as exc:
|
||||
error_text = _short(str(exc), 800)
|
||||
attempts.append(f"{model}: invalid report ({_short(str(exc), 160)})")
|
||||
log.warning(
|
||||
"model %s returned invalid report, retrying once with repair hint: %s",
|
||||
model,
|
||||
_short(str(exc), 300),
|
||||
)
|
||||
|
||||
repair_messages = [
|
||||
*messages,
|
||||
{"role": "assistant", "content": _short(content, 12000)},
|
||||
{"role": "user", "content": REPAIR_INSTRUCTION + error_text},
|
||||
]
|
||||
content = self._call(model, repair_messages, attempts, label="repair")
|
||||
if content is None:
|
||||
return None
|
||||
try:
|
||||
return parse_report(content)
|
||||
except (ValidationError, ValueError) as exc:
|
||||
attempts.append(f"{model} (repair): invalid report ({_short(str(exc), 160)})")
|
||||
log.warning("model %s repair retry still invalid: %s", model, _short(str(exc), 300))
|
||||
return None
|
||||
|
||||
def _call(
|
||||
self, model: str, messages: list[dict[str, Any]], attempts: list[str], label: str = "initial"
|
||||
) -> str | None:
|
||||
"""Perform one chat completion. Returns content text, or None (attempt recorded) on failure."""
|
||||
tag = model if label == "initial" else f"{model} ({label})"
|
||||
extra_body: dict[str, Any] = {"reasoning_split": True}
|
||||
if model.startswith("MiniMax-M3"):
|
||||
extra_body["thinking"] = {"type": "adaptive"}
|
||||
try:
|
||||
completion = self._client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
max_completion_tokens=self._settings.minimax_max_completion_tokens,
|
||||
response_format=self._response_format,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
except OPENAI_ERRORS as exc:
|
||||
status = getattr(exc, "status_code", None)
|
||||
detail = f"{type(exc).__name__}" + (f" {status}" if status else "") + f": {_short(str(exc), 160)}"
|
||||
attempts.append(f"{tag}: {detail}")
|
||||
log.warning("MiniMax call failed for %s: %s", tag, detail)
|
||||
return None
|
||||
|
||||
choices = getattr(completion, "choices", None) or []
|
||||
if not choices:
|
||||
attempts.append(f"{tag}: empty choices")
|
||||
log.warning("MiniMax %s returned no choices", tag)
|
||||
return None
|
||||
choice = choices[0]
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
self._log_usage(model, completion, finish_reason)
|
||||
if finish_reason == "length":
|
||||
attempts.append(f"{tag}: truncated (finish_reason=length)")
|
||||
log.warning("MiniMax %s output truncated (finish_reason=length)", tag)
|
||||
return None
|
||||
content = getattr(choice.message, "content", None)
|
||||
if not content:
|
||||
attempts.append(f"{tag}: empty content")
|
||||
log.warning("MiniMax %s returned empty content", tag)
|
||||
return None
|
||||
return content
|
||||
|
||||
@staticmethod
|
||||
def _log_usage(model: str, completion: Any, finish_reason: str | None) -> None:
|
||||
usage = getattr(completion, "usage", None)
|
||||
prompt_tokens = getattr(usage, "prompt_tokens", None)
|
||||
completion_tokens = getattr(usage, "completion_tokens", None)
|
||||
details = getattr(usage, "prompt_tokens_details", None)
|
||||
cached = getattr(details, "cached_tokens", None) if details is not None else None
|
||||
log.info(
|
||||
"MiniMax response model=%s served_as=%s finish=%s prompt_tokens=%s completion_tokens=%s "
|
||||
"cached_tokens=%s",
|
||||
model,
|
||||
getattr(completion, "model", None),
|
||||
finish_reason,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
cached,
|
||||
)
|
||||
|
||||
|
||||
def _short(text: str, limit: int) -> str:
|
||||
text = " ".join(text.split())
|
||||
return text if len(text) <= limit else text[: limit - 1] + "…"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"OPENAI_ERRORS",
|
||||
"MiniMaxProvider",
|
||||
"coerce_report_data",
|
||||
"extract_json_text",
|
||||
"parse_report",
|
||||
"report_json_schema",
|
||||
"report_response_format",
|
||||
]
|
||||
Reference in New Issue
Block a user