memory-system/lifecycle/llm_parse.py

"""Shared LLM response parsing utilities for lifecycle modules."""
import json
import logging
import re
import time
from typing import Optional

logger = logging.getLogger(__name__)

MAX_PARSE_RETRIES = 2  # 3 total attempts (1 initial + 2 retry)

RETRY_SYSTEM_PROMPT = (
    '只输出原始JSON。不要有任何其他文字、分析、思考、解释、markdown。\n'
    '输出格式示例：\n'
    '[{"name": "场景名", "summary": "摘要", "memory_ids": ["id1", "id2"]}]\n'
    '重要：你回复的第一个字符必须是 [，最后一个字符必须是 ]。'
)

RETRY_OBJECT_PROMPT = (
    '只输出原始JSON对象。不要有任何其他文字、分析、思考、解释、markdown。\n'
    '输出格式：\n'
    '{"data": [{"name": "场景名", "summary": "摘要", "memory_ids": ["id1", "id2"]}]}\n'
    '重要：你回复的第一个字符必须是 {，最后一个字符必须是 }。'
)


def _strip_markdown_code_block(text: str) -> str:
    text = text.strip()
    if not text.startswith("```"):
        return text
    text = text.split("\n", 1)[1] if "\n" in text else text[3:]
    if text.endswith("```"):
        text = text[:-3].strip()
    if text.lower().startswith("json"):
        text = text[4:].strip()
    elif text.lower().startswith("json\n"):
        text = text[5:].strip()
    return text.strip()


def _strip_reasoning_prefix(text: str) -> str:
    """Strip reasoning/analysis text that some models output before JSON.

    Models like deepseek-v4-flash often output chain-of-thought analysis
    before the actual JSON. This function tries to find where the JSON
    actually starts by looking for lines that start with [ or { after
    stripping analysis text.
    """
    text = text.strip()

    # If text starts with [ or { it might already be JSON
    if text and text[0] in ('[', '{'):
        return text

    # Try to find the last occurrence of a JSON-like pattern
    lines = text.split('\n')
    for i in range(len(lines) - 1, -1, -1):
        stripped = lines[i].strip()
        if stripped.startswith('[') or stripped.startswith('{'):
            candidate = '\n'.join(lines[i:])
            try:
                json.loads(candidate)
                return candidate
            except (json.JSONDecodeError, ValueError):
                continue

    # Try finding first [ or { that starts a valid JSON
    for start_char in ['[', '{']:
        idx = text.find(start_char)
        while idx >= 0:
            candidate = text[idx:]
            end_char = ']' if start_char == '[' else '}'
            if candidate.rstrip().endswith(end_char):
                try:
                    json.loads(candidate)
                    return candidate
                except (json.JSONDecodeError, ValueError):
                    pass
            idx = text.find(start_char, idx + 1)

    return text


def _find_json_string(text: str) -> Optional[str]:
    """Find the first valid JSON array or object in text.

    Iterates through ALL bracket-delimited segments, not just from the
    first bracket. This handles cases where the LLM outputs analysis
    text containing bracketed IDs before the actual JSON.
    """
    text = text.strip()
    if not text:
        return None

    # Fast path: entire text is valid JSON
    if (text.startswith("[") or text.startswith("{")) and text[-1] in ("}", "]"):
        try:
            json.loads(text)
            return text
        except (json.JSONDecodeError, ValueError):
            pass

    # Strip markdown code blocks
    cleaned = _strip_markdown_code_block(text)
    if cleaned != text:
        return _find_json_string(cleaned)

    # Strip reasoning prefix
    cleaned = _strip_reasoning_prefix(text)
    if cleaned != text:
        result = _find_json_string(cleaned)
        if result:
            return result

    # Scan all bracket-delimited segments (for "[" / "]" and "{" / "}")
    for start_char, end_char in [("[", "]"), ("{", "}")]:
        search_start = 0
        while True:
            start_idx = text.find(start_char, search_start)
            if start_idx == -1:
                break
            depth = 0
            for i in range(start_idx, len(text)):
                if text[i] == start_char:
                    depth += 1
                elif text[i] == end_char:
                    depth -= 1
                    if depth == 0:
                        candidate = text[start_idx:i + 1]
                        try:
                            json.loads(candidate)
                            return candidate
                        except (json.JSONDecodeError, ValueError):
                            # Move past this bracket pair and continue scanning
                            search_start = i + 1
                            break
            else:
                # No matching end bracket found
                search_start = start_idx + 1

    # Regex fallback: find ALL bracket-delimited segments
    for pattern in [r"\[.*\]", r"\{.*\}"]:
        for match in re.finditer(pattern, text, re.DOTALL):
            candidate = match.group(0)
            try:
                json.loads(candidate)
                return candidate
            except (json.JSONDecodeError, ValueError):
                continue

    return None


def _try_fix_json(text: str) -> Optional[str]:
    """Try to fix common JSON issues from LLM output."""
    # Remove trailing commas before } or ]
    text = re.sub(r",\s*}", "}", text)
    text = re.sub(r",\s*\]", "]", text)
    # Replace single quotes with double quotes only if no double quotes present
    if '"' not in text:
        text = text.replace("'", '"')
    # Remove BOM and zero-width characters
    text = text.replace("\ufeff", "").replace("\u200b", "")
    # Try parsing as-is first
    try:
        json.loads(text)
        return text
    except json.JSONDecodeError:
        pass
    # Try fixing unescaped newlines inside JSON string values
    fixed = re.sub(r'(?<=: ")((?:[^"]|\\")*?)(\n)((?:[^"]|\\")*?(?="))', r'\1\\n\3', text)
    try:
        json.loads(fixed)
        return fixed
    except json.JSONDecodeError:
        pass
    # Try wrapping in array if it looks like a bare object
    stripped = text.strip()
    if stripped.startswith("{") and not stripped.startswith("["):
        wrapped = "[" + stripped + "]"
        try:
            json.loads(wrapped)
            return wrapped
        except json.JSONDecodeError:
            pass
    # Try to extract JSON from truncated output: find last complete object in array
    if stripped.startswith("["):
        depth = 0
        objects = []
        obj_start = -1
        for i, ch in enumerate(stripped):
            if ch == "{":
                if depth == 0:
                    obj_start = i
                depth += 1
            elif ch == "}":
                depth -= 1
                if depth == 0 and obj_start >= 0:
                    objects.append(stripped[obj_start:i + 1])
                    obj_start = -1
        if objects:
            candidate = "[" + ",".join(objects) + "]"
            try:
                json.loads(candidate)
                return candidate
            except json.JSONDecodeError:
                pass
    return None


def parse_llm_json_array(llm_client, system_prompt, user_prompt,
                         temperature=0.3, context="", max_tokens=4000):
    last_raw = ""
    total = 1 + MAX_PARSE_RETRIES
    t_start = time.time()
    for attempt in range(total):
        is_retry = attempt > 0

        if is_retry:
            if attempt >= total - 1:
                # Final attempt: json_mode as last resort with wrapper object
                prompt = RETRY_OBJECT_PROMPT
                use_json_mode = True
                prefill_val = None
            else:
                # Intermediate retry: use prefill "["
                prompt = RETRY_SYSTEM_PROMPT
                use_json_mode = False
                prefill_val = "["
        else:
            # First attempt: prefill "[" to force JSON output immediately
            prompt = system_prompt
            use_json_mode = False
            prefill_val = "["

        response = llm_client.chat(
            prompt, user_prompt,
            temperature=temperature, json_mode=use_json_mode,
            max_tokens=max_tokens,
            thinking={"type": "disabled"},
            prefill=prefill_val,
        )
        if not response:
            logger.warning("[%s] Attempt %d: LLM returned empty", context, attempt + 1)
            continue
        last_raw = response
        json_str = _find_json_string(response)
        if json_str:
            try:
                result = json.loads(json_str)
            except json.JSONDecodeError as e:
                logger.warning("[%s] Attempt %d: JSON parse error: %s", context, attempt + 1, e)
                continue
            if isinstance(result, list):
                elapsed = time.time() - t_start
                logger.info("[%s] Parsed JSON array in %.1fs (%d attempt(s), %d items)", context, elapsed, attempt + 1, len(result))
                return result
            if isinstance(result, dict):
                if "data" in result and isinstance(result["data"], list):
                    elapsed = time.time() - t_start
                    logger.info("[%s] Parsed JSON array in %.1fs (%d attempt(s), %d items)", context, elapsed, attempt + 1, len(result["data"]))
                    return result["data"]
                logger.warning("[%s] Attempt %d: LLM returned dict, wrapping in list", context, attempt + 1)
                return [result]
        fixed = _try_fix_json(response)
        if fixed:
            try:
                result = json.loads(fixed)
            except json.JSONDecodeError:
                continue
            if isinstance(result, list):
                return result
            if isinstance(result, dict):
                if "data" in result and isinstance(result["data"], list):
                    return result["data"]
                return [result]
        logger.warning(
            "[%s] Attempt %d/%d: No valid JSON array found (len=%d): %r",
            context, attempt + 1, total, len(response), response[:500]
        )
    logger.error("[%s] All %d attempts failed. Last raw: %r", context, total, last_raw[:500])
    raise ValueError("LLM failed to return valid JSON after retries")


def parse_llm_json_object(llm_client, system_prompt, user_prompt,
                          temperature=0.3, context="", max_tokens=4000):
    last_raw = ""
    total = 1 + MAX_PARSE_RETRIES
    t_start = time.time()
    for attempt in range(total):
        is_retry = attempt > 0

        if is_retry:
            if attempt >= total - 1:
                # Final attempt: json_mode as last resort
                prompt = '只输出 JSON 对象。不要输出任何其他文字。\n输出格式：{"preferences":[],"habits":[],"expertise":[],"communication_style":"","summary":""}'
                use_json_mode = True
                prefill_val = None
            else:
                prompt = '只输出 JSON 对象。不要输出任何其他文字。\n禁止输出：分析、思考、解释、讨论、描述、markdown。'
                use_json_mode = False
                prefill_val = "{"
        else:
            # First attempt: prefill "{" to force JSON output
            prompt = system_prompt
            use_json_mode = False
            prefill_val = "{"

        response = llm_client.chat(
            prompt, user_prompt,
            temperature=temperature, json_mode=use_json_mode,
            max_tokens=max_tokens,
            thinking={"type": "disabled"},
            prefill=prefill_val,
        )
        if not response:
            logger.warning("[%s] Attempt %d: LLM returned empty", context, attempt + 1)
            continue
        last_raw = response
        json_str = _find_json_string(response)
        if json_str:
            try:
                result = json.loads(json_str)
            except json.JSONDecodeError as e:
                logger.warning("[%s] Attempt %d: JSON parse error: %s", context, attempt + 1, e)
                continue
            if isinstance(result, dict):
                elapsed = time.time() - t_start
                logger.info("[%s] Parsed JSON object in %.1fs (%d attempt(s))", context, elapsed, attempt + 1)
                return result
        fixed = _try_fix_json(response)
        if fixed:
            try:
                result = json.loads(fixed)
            except json.JSONDecodeError:
                continue
            if isinstance(result, dict):
                return result
        logger.warning(
            "[%s] Attempt %d/%d: No valid JSON object found (len=%d): %r",
            context, attempt + 1, total, len(response), response[:500]
        )
    logger.error("[%s] All %d attempts failed. Last raw: %r", context, total, last_raw[:500])
    raise ValueError("LLM failed to return valid JSON after retries")