Fix memoir image prompt parsing and host allowlist

This commit is contained in:
Kevin
2026-03-11 13:18:20 +08:00
parent 1f98b8bfd6
commit 32954d4b3f
5 changed files with 100 additions and 3 deletions

View File

@@ -7,6 +7,7 @@ from .settings import MemoirImageSettings
logger = logging.getLogger(__name__)
_CJK_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff]")
_MARKDOWN_JSON_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*(.*?)\s*```\s*$", re.IGNORECASE | re.DOTALL)
class MemoirImagePromptService:
@@ -54,13 +55,15 @@ class MemoirImagePromptService:
}
if self.llm:
raw_response = None
try:
response = self.llm.invoke(
"Return JSON only with keys prompt, style, size. "
"Convert the memoir scene into an image-generation prompt.\n"
+ json.dumps(llm_input, ensure_ascii=False)
)
parsed = json.loads(response.content)
raw_response = response.content
parsed = json.loads(_extract_json_payload(raw_response))
return {
"prompt": _ensure_style_in_prompt(parsed["prompt"], parsed.get("style", style)),
"style": parsed.get("style", style),
@@ -118,6 +121,23 @@ def _contains_cjk(value: str) -> bool:
return bool(_CJK_RE.search(value or ""))
def _extract_json_payload(raw_response: str | None) -> str:
cleaned = (raw_response or "").strip()
fenced_match = _MARKDOWN_JSON_FENCE_RE.match(cleaned)
if fenced_match:
cleaned = fenced_match.group(1).strip()
if cleaned.startswith("{") and cleaned.endswith("}"):
return cleaned
start = cleaned.find("{")
end = cleaned.rfind("}")
if start != -1 and end != -1 and end > start:
return cleaned[start : end + 1].strip()
return cleaned
def _ensure_style_in_prompt(prompt: str, style: str) -> str:
cleaned_prompt = (prompt or "").strip()
cleaned_style = (style or "").strip()