claude-smart 0.2.46 → 0.2.48

package/plugin/vendor/reflexio/reflexio/server/llm/_litellm_structured_output.py

@@ -0,0 +1,195 @@
+"""Structured-output concern for ``LiteLLMClient`` — schema selection + parsing (Tier-2.5).
+
+``StructuredOutputMixin`` holds the response-format selection helpers
+(``_supports_response_schema``, ``_provider_for_model``,
+``_accepts_json_schema_response_format``, ``_provider_response_format``) and the
+post-hoc parse orchestrator (``_maybe_parse_structured_output``), plus the
+``_JSON_SCHEMA_PROVIDER_ALLOWLIST`` class constant they use.
+
+SINK-1 (patch-where-used): ``assert_provider_safe_schema`` is imported HERE, so
+tests patch it at ``_litellm_structured_output.assert_provider_safe_schema``
+(the real guard raises under pytest; patching the old facade namespace would
+no-op). ``litellm.supports_response_schema``/``get_llm_provider`` are called via
+the shared ``litellm`` module attr.
+
+Bodies moved VERBATIM; only the ``self``-typing (per-mixin TYPE_CHECKING stub of
+``config``, Tier-1b idiom) is added. ``_provider_response_format`` and
+``_maybe_parse_structured_output`` are the cross-mixin edges text-gen depends on,
+so this module moves BEFORE ``_litellm_text_generation``.
+"""
+
+import json
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any
+
+import litellm
+from pydantic import BaseModel
+
+from reflexio.server.llm._litellm_json_extraction import (
+    _extract_json_from_string,
+    _looks_truncated_json,
+    _sanitize_json_string,
+)
+from reflexio.server.llm._litellm_types import StructuredOutputParseError
+from reflexio.server.llm.llm_utils import (
+    assert_provider_safe_schema,
+    is_pydantic_model,
+    strict_response_format_for_model,
+)
+
+if TYPE_CHECKING:
+    from reflexio.server.llm._litellm_types import LiteLLMConfig
+
+
+class StructuredOutputMixin:
+    """Response-format schema selection + structured-output parsing.
+
+    Mixed into ``LiteLLMClient``; ``self.config`` is owned by the client-core
+    ``__init__`` on the facade. The annotation-only stub below (Tier-1b idiom)
+    gives pyright the foreign-member type without shared class-level state.
+    """
+
+    # OpenAI-compatible providers that accept a ``json_schema`` response_format
+    # but that ``litellm.supports_response_schema`` reports as unsupported. For
+    # these, the gate below would fall back to handing LiteLLM the raw Pydantic
+    # model; LiteLLM then builds the ``json_schema`` itself and emits ``oneOf``
+    # for discriminated unions, which strict structured-output endpoints reject
+    # (Sentry PYTHON-FASTAPI-9J). Listing the provider here forces our own
+    # normalized strict schema (``oneOf`` folded into ``anyOf``) to be sent.
+    _JSON_SCHEMA_PROVIDER_ALLOWLIST: frozenset[str] = frozenset({"minimax"})
+
+    # Base-owned attribute read for the parse-failure error message (init'd in
+    # the facade ``__init__``). Annotation-only; NEVER assign here.
+    config: "LiteLLMConfig"
+
+    @staticmethod
+    @lru_cache(maxsize=256)
+    def _supports_response_schema(model: str) -> bool:
+        try:
+            return bool(litellm.supports_response_schema(model=model))
+        except Exception:
+            return False
+
+    @staticmethod
+    @lru_cache(maxsize=256)
+    def _provider_for_model(model: str) -> str | None:
+        try:
+            return litellm.get_llm_provider(model)[1]
+        except Exception:
+            return None
+
+    @classmethod
+    def _accepts_json_schema_response_format(cls, model: str) -> bool:
+        """Whether to send ``model`` an explicit strict ``json_schema`` schema.
+
+        True when LiteLLM reports native response-schema support, or when the
+        provider is a known OpenAI-compatible endpoint that LiteLLM
+        under-reports (see ``_JSON_SCHEMA_PROVIDER_ALLOWLIST``). In the latter
+        case LiteLLM would otherwise forward a ``json_schema`` it built itself,
+        emitting ``oneOf`` for discriminated unions that the endpoint rejects.
+        """
+        if cls._supports_response_schema(model):
+            return True
+        return cls._provider_for_model(model) in cls._JSON_SCHEMA_PROVIDER_ALLOWLIST
+
+    def _provider_response_format(
+        self,
+        *,
+        response_format: Any,
+        model: str,
+        strict_response_format: bool,
+    ) -> Any:
+        """Return the provider-facing response_format while preserving parser schema.
+
+        Callers pass a Pydantic model so local parsing stays type-safe. When the
+        target model accepts a JSON Schema response format — either LiteLLM
+        reports native support, or the provider is an OpenAI-compatible endpoint
+        LiteLLM under-reports (see ``_accepts_json_schema_response_format``) — we
+        send an explicit strict schema to constrain generation. Truly
+        unsupported providers keep the existing Pydantic response_format
+        behavior.
+        """
+
+        if not is_pydantic_model(response_format):
+            return response_format
+
+        # Build the native schema once and reuse it for both the boundary guard and
+        # (when applicable) the strict normalizer, avoiding a second schema build.
+        # Boundary guard: models inheriting StrictStructuredOutput are safe by
+        # construction; this catches a model that forgot the base (raises under
+        # tests, warns in prod) regardless of which path is taken below.
+        schema = response_format.model_json_schema()
+        assert_provider_safe_schema(schema, name=response_format.__name__)
+
+        if strict_response_format and self._accepts_json_schema_response_format(model):
+            return strict_response_format_for_model(response_format, schema=schema)
+        return response_format
+
+    def _maybe_parse_structured_output(
+        self,
+        content: str,
+        response_format: Any,
+        parse_structured_output: bool,
+    ) -> str | BaseModel:
+        """
+        Parse structured output if applicable.
+
+        Args:
+            content: Raw response content.
+            response_format: Expected response format (must be a Pydantic BaseModel class).
+            parse_structured_output: Whether to parse the output.
+
+        Returns:
+            String for text responses, or BaseModel instance for structured responses.
+        """
+        if not response_format or not parse_structured_output:
+            return content
+
+        if content is None:
+            return content
+
+        # If content is already a Pydantic model (some providers return parsed)
+        if isinstance(content, BaseModel):
+            return content
+
+        # Try to parse JSON and convert to Pydantic model
+        # Extract JSON from markdown code blocks if present
+        json_str = _extract_json_from_string(content)
+        try:
+            parsed = json.loads(json_str)
+
+            # response_format must be a Pydantic model (validated at entry points)
+            return response_format.model_validate(parsed)
+        except Exception:
+            # LLMs sometimes produce Python-style output (single quotes, True/False,
+            # trailing commas). Try to sanitize before giving up.
+            try:
+                sanitized = _sanitize_json_string(json_str)
+                parsed = json.loads(sanitized)
+                return response_format.model_validate(parsed)
+            except Exception:
+                # Last resort: json-repair can recover complete responses with
+                # small syntax glitches, such as missing commas. Do not repair
+                # likely truncation: the retry loop should request a fresh
+                # complete response instead of accepting invented tail content.
+                try:
+                    from json_repair import repair_json
+
+                    if _looks_truncated_json(json_str):
+                        raise StructuredOutputParseError(
+                            "Structured output appears truncated"
+                        )
+
+                    repaired = repair_json(json_str, return_objects=True)
+                    return response_format.model_validate(repaired)
+                except Exception as e:
+                    model = self.config.model
+                    snippet = (
+                        content[:200]
+                        if isinstance(content, str)
+                        else repr(content)[:200]
+                    )
+                    raise StructuredOutputParseError(
+                        f"Structured output parse failed for model={model!r}: {e}. "
+                        f"Content snippet: {snippet!r}"
+                    ) from e

package/plugin/vendor/reflexio/reflexio/server/llm/_litellm_subprocess.py

@@ -0,0 +1,152 @@
+"""Subprocess-isolation helpers for the hard-timeout completion path (Tier-2.5 leaf).
+
+Stateless leaf module: the picklable ``_Completion*Snapshot`` dataclasses plus the
+snapshot builders and the ``_litellm_completion_worker`` that runs
+``litellm.completion`` inside a child process. Used by ``TextGenerationMixin``'s
+``_completion_with_hard_timeout`` (Task 4) which sizes and kills the subprocess.
+
+LLM-mock: the worker calls ``litellm.completion`` via the MODULE ATTR (never
+``from litellm import completion``) so the global ``patch("litellm.completion")``
+mock is inherited across the ``fork`` start-method. The multiprocessing
+start-method is intentionally left unchanged.
+
+SINK-2 (identity): the 6 snapshot classes and ``_litellm_completion_worker`` are
+re-exported by import binding from the facade — they must be the SAME class/object
+the worker constructs and tests ``isinstance``-check. Bodies moved VERBATIM.
+"""
+
+import multiprocessing
+import pickle
+from dataclasses import dataclass, field
+from typing import Any
+
+import litellm
+
+
+@dataclass
+class _CompletionMessageSnapshot:
+    content: str | None = None
+    tool_calls: Any | None = None
+
+
+@dataclass
+class _CompletionChoiceSnapshot:
+    message: _CompletionMessageSnapshot
+    finish_reason: str | None = None
+
+
+@dataclass
+class _PromptTokenDetailsSnapshot:
+    cached_tokens: int = 0
+
+
+@dataclass
+class _CompletionUsageSnapshot:
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+    total_tokens: int | None = None
+    prompt_tokens_details: _PromptTokenDetailsSnapshot | None = None
+    cache_creation_input_tokens: int | None = None
+    cache_read_input_tokens: int | None = None
+
+
+@dataclass
+class _CompletionResponseSnapshot:
+    choices: list[_CompletionChoiceSnapshot]
+    usage: _CompletionUsageSnapshot | None = None
+    model: str | None = None
+    _hidden_params: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class _CompletionErrorSnapshot:
+    type_name: str
+    message: str
+    model: str | None = None
+    llm_provider: str | None = None
+
+
+def _snapshot_completion_error(
+    exc: BaseException, params: dict[str, Any]
+) -> _CompletionErrorSnapshot:
+    model = getattr(exc, "model", None) or params.get("model")
+    llm_provider = getattr(exc, "llm_provider", None)
+    return _CompletionErrorSnapshot(
+        type_name=type(exc).__name__,
+        message=str(exc),
+        model=str(model) if model else None,
+        llm_provider=str(llm_provider) if llm_provider else None,
+    )
+
+
+def _ensure_picklable(value: Any) -> Any:
+    try:
+        pickle.dumps(value)
+    except Exception:
+        return repr(value)
+    return value
+
+
+def _snapshot_completion_response(response: Any) -> _CompletionResponseSnapshot:
+    choices: list[_CompletionChoiceSnapshot] = []
+    for choice in getattr(response, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        choices.append(
+            _CompletionChoiceSnapshot(
+                message=_CompletionMessageSnapshot(
+                    content=getattr(message, "content", None),
+                    tool_calls=_ensure_picklable(getattr(message, "tool_calls", None)),
+                ),
+                finish_reason=getattr(choice, "finish_reason", None),
+            )
+        )
+
+    usage = getattr(response, "usage", None)
+    usage_snapshot = None
+    if usage is not None:
+        prompt_details = getattr(usage, "prompt_tokens_details", None)
+        prompt_details_snapshot = None
+        if prompt_details is not None:
+            prompt_details_snapshot = _PromptTokenDetailsSnapshot(
+                cached_tokens=int(getattr(prompt_details, "cached_tokens", 0) or 0)
+            )
+        usage_snapshot = _CompletionUsageSnapshot(
+            prompt_tokens=getattr(usage, "prompt_tokens", None),
+            completion_tokens=getattr(usage, "completion_tokens", None),
+            total_tokens=getattr(usage, "total_tokens", None),
+            prompt_tokens_details=prompt_details_snapshot,
+            cache_creation_input_tokens=getattr(
+                usage, "cache_creation_input_tokens", None
+            ),
+            cache_read_input_tokens=getattr(usage, "cache_read_input_tokens", None),
+        )
+
+    hidden_params = getattr(response, "_hidden_params", {}) or {}
+    if not isinstance(hidden_params, dict):
+        hidden_params = {}
+
+    return _CompletionResponseSnapshot(
+        choices=choices,
+        usage=usage_snapshot,
+        model=getattr(response, "model", None),
+        _hidden_params={str(k): _ensure_picklable(v) for k, v in hidden_params.items()},
+    )
+
+
+def _picklable_completion_result(response: Any) -> Any:
+    try:
+        pickle.dumps(response)
+    except Exception:
+        return _snapshot_completion_response(response)
+    return response
+
+
+def _litellm_completion_worker(
+    params: dict[str, Any], result_queue: multiprocessing.Queue
+) -> None:
+    try:
+        result_queue.put(
+            ("ok", _picklable_completion_result(litellm.completion(**params)))
+        )
+    except BaseException as exc:
+        result_queue.put(("error", _snapshot_completion_error(exc, params)))