main/utils.py

"""Utility functions for the inference API."""
import json
import re
from typing import Dict, Any

def extract_json(text: str) -> Dict[str, Any]:
    """Extract JSON from text that might contain other content.

    Handles cases like:
    - Clean JSON: {"key": "value"}
    - JSON with prefix: Sure! Here's your JSON: {"key": "value"}
    - JSON with suffix: {"key": "value"} Let me know if you need anything else!
    """
    # Find anything that looks like a JSON object
    json_pattern = r'\{(?:[^{}]|(?R))*\}'
    matches = re.finditer(json_pattern, text)

    # Try each match until we find valid JSON
    for match in matches:
        try:
            potential_json = match.group()
            parsed = json.loads(potential_json)
            return parsed
        except json.JSONDecodeError:
            continue

    # If we couldn't find any valid JSON, raise an error
    raise ValueError("No valid JSON found in response")