llm_json.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Parse JSON objects from LLM chat responses."""
from __future__ import annotations
 
import json
import re
from typing import Any
 
 
def parse_llm_json_object(raw: str) -> dict[str, Any]:
    """Extract the first JSON object from an LLM reply (handles fences and trailing text)."""
    text = (raw or "").strip()
    if not text:
        raise ValueError("empty response")
 
    if text.startswith("```"):
        parts = text.split("```")
        if len(parts) >= 2:
            text = parts[1].strip()
            if text.lower().startswith("json"):
                text = text[4:].lstrip()
 
    start = text.find("{")
    if start == -1:
        raise ValueError("no JSON object found")
    text = text[start:]
 
    try:
        obj, _ = json.JSONDecoder().raw_decode(text)
    except json.JSONDecodeError:
        match = re.search(r"\{[\s\S]*\}", text)
        if not match:
            raise
        obj, _ = json.JSONDecoder().raw_decode(match.group(0))
 
    if not isinstance(obj, dict):
        raise ValueError("expected JSON object")
    return obj
 
 
def sentence_lists_from_llm(raw: str) -> tuple[list[str], list[str]]:
    """Parse positive/negative sentence lists from an LLM JSON reply."""
    data = parse_llm_json_object(raw)
    positive = _as_str_list(data.get("positive"))
    negative = _as_str_list(data.get("negative"))
    if not positive and not negative:
        raise ValueError("empty sentence lists")
    return positive, negative
 
 
def _as_str_list(value: Any) -> list[str]:
    if not isinstance(value, list):
        return []
    return [str(item).strip() for item in value if str(item).strip()]
 
 
def llm_sentence_generation_error() -> str:
    """User-facing message when LLM example generation fails."""
    return (
        "Could not generate example sentences for this feature. "
        "Try again in a moment, or use a larger --corpus."
    )
 
 
def chat_json_completion(client: Any, **kwargs: Any) -> Any:
    """Chat completion preferring JSON object mode, with a plain fallback."""
    try:
        return client.chat.completions.create(
            **kwargs,
            response_format={"type": "json_object"},
        )
    except Exception:
        kwargs.pop("response_format", None)
        return client.chat.completions.create(**kwargs)