PMFlow/backend/app/services/llm_service/mixin_structured_parse.py

from __future__ import annotations

from .shared import *

class LLMStructuredParseMixin:
    @staticmethod
    def _extract_json_candidate(text: str) -> str:
        raw = str(text or "").strip()
        if not raw:
            return ""
        if raw.startswith("```"):
            chunks = raw.split("```")
            if len(chunks) >= 2:
                raw = chunks[1].strip()
            if raw.startswith("json"):
                raw = raw[4:].strip()
        if raw.startswith("{") or raw.startswith("["):
            return raw

        object_candidate = LLMStructuredParseMixin._extract_balanced_json_segment(raw, "{", "}")
        array_candidate = LLMStructuredParseMixin._extract_balanced_json_segment(raw, "[", "]")
        if object_candidate and array_candidate:
            return object_candidate if raw.find("{") <= raw.find("[") else array_candidate
        if object_candidate:
            return object_candidate
        if array_candidate:
            return array_candidate
        return raw

    @staticmethod
    def _looks_like_json_container(text: str) -> bool:
        raw = str(text or "").strip()
        if not raw:
            return False
        if raw.startswith("{") or raw.startswith("["):
            return True
        return "{" in raw or "[" in raw

    @staticmethod
    def _extract_balanced_json_segment(text: str, open_char: str, close_char: str) -> str:
        raw = str(text or "")
        start = raw.find(open_char)
        if start < 0:
            return ""
        depth = 0
        in_str = False
        escaped = False
        for idx in range(start, len(raw)):
            ch = raw[idx]
            if in_str:
                if escaped:
                    escaped = False
                elif ch == "\\":
                    escaped = True
                elif ch == '"':
                    in_str = False
                continue
            if ch == '"':
                in_str = True
                continue
            if ch == open_char:
                depth += 1
            elif ch == close_char:
                depth -= 1
                if depth == 0:
                    return raw[start : idx + 1].strip()
        return ""

    def _is_invalid_structured_output(self, *, function_name: str, response_text: str) -> bool:
        if not self._structured_output_required(function_name):
            return False
        candidate = self._extract_json_candidate(response_text)
        if not candidate:
            return True
        try:
            parsed = json.loads(candidate)
        except Exception:
            return True
        parsed = self._unwrap_structured_payload_once(parsed)
        parsed = self._coerce_structured_payload(
            function_name=function_name,
            payload=parsed,
        )
        validation_error = self._validate_structured_payload(function_name=function_name, payload=parsed)
        if validation_error is not None:
            return True
        return not isinstance(parsed, (dict, list))

    @staticmethod
    def _strip_code_fence(text: str) -> str:
        raw = str(text or "").strip()
        if not raw.startswith("```"):
            return raw
        chunks = raw.split("```")
        if len(chunks) >= 2:
            raw = chunks[1].strip()
        if raw.lower().startswith("json"):
            raw = raw[4:].strip()
        return raw.strip()

    def _extract_text_from_chat_response(
        self,
        *,
        message: Any,
        choice: Any,
        response: Any,
        structured_output: bool,
    ) -> str:
        if structured_output:
            for _, raw_value in self._iter_structured_candidates(
                message=message,
                choice=choice,
                response=response,
                include_reasoning=False,
            ):
                text = self._sanitize_visible_answer(self._extract_text_from_value(raw_value))
                normalized = self._strip_code_fence(text)
                if not self._looks_like_json_container(normalized):
                    continue
                candidate = self._extract_json_candidate(normalized)
                if candidate:
                    return candidate
            return ""
        if isinstance(message, dict):
            text = self._sanitize_visible_answer(self._extract_text_from_value(message.get("content")))
            if text:
                return text
            tool_calls = message.get("tool_calls")
            if isinstance(tool_calls, list) and tool_calls:
                first = tool_calls[0] if isinstance(tool_calls[0], dict) else {}
                fn = first.get("function") if isinstance(first, dict) else {}
                if isinstance(fn, dict):
                    text = self._sanitize_visible_answer(self._extract_text_from_value(fn.get("arguments")))
                    if text:
                        return text
        if isinstance(choice, dict):
            text = self._sanitize_visible_answer(self._extract_text_from_value(choice.get("text")))
            if text:
                return text
        if isinstance(response, dict):
            text = self._sanitize_visible_answer(self._extract_text_from_value(response.get("output_text")))
            if text:
                return text
        return ""

    def _extract_reasoning_fallback_text(
        self,
        *,
        message: Any,
        choice: Any,
        response: Any,
    ) -> str:
        candidates: list[Any] = []
        if isinstance(message, dict):
            candidates.append(message.get("reasoning_content"))
            candidates.append(message.get("reasoning"))
        if isinstance(choice, dict):
            candidates.append(choice.get("reasoning_content"))
            candidates.append(choice.get("reasoning"))
        if isinstance(response, dict):
            candidates.append(response.get("reasoning_content"))
            candidates.append(response.get("reasoning"))
        for raw in candidates:
            text = self._sanitize_visible_answer(self._extract_text_from_value(raw))
            if text:
                return text
        return ""

    @staticmethod
    def _reasoning_content_fallback_enabled(*, function_name: str) -> bool:
        fn = str(function_name or "").strip().lower()
        # P0 guardrail: user-facing answer/reply/qa must never expose reasoning_content.
        if fn and any(token in fn for token in ("answer", "reply", "qa")):
            return False
        if fn and "title_refine" in fn:
            return True
        raw = str(os.getenv("LLM_REASONING_CONTENT_FALLBACK", "0")).strip().lower()
        if raw not in {"1", "true", "yes", "on"}:
            return False
        if not fn:
            return True
        if fn.startswith("generate_structured") or fn.startswith("llm_generate_structured"):
            return False
        if "signal_batch_analysis" in fn:
            return False
        return False