main/utils.py

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

import yaml


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")

def load_config():
    """
    Load configuration from config files in the resources directory.
    Uses CONFIG_ENV environment variable to determine which config to load.
    Defaults to 'local' if no environment is specified.
    """
    # Get environment name from env var, default to 'local'
    env_name = os.environ.get("CONFIG_ENV", "local")

    # Construct path to resources directory and config file
    resources_dir = Path(__file__).parent / "resources"
    config_path = resources_dir / f"{env_name}_config.yaml"

    # Create resources directory if it doesn't exist
    resources_dir.mkdir(exist_ok=True)

    # Check if config file exists
    if not config_path.exists():
        logging.warning(f"Config file {config_path} not found, falling back to local_config.yaml")
        config_path = resources_dir / "local_config.yaml"

        # If even local config doesn't exist, raise error
        if not config_path.exists():
            raise FileNotFoundError(
                f"No configuration file found at {config_path}. "
                "Please ensure at least local_config.yaml exists in the resources directory."
            )

    # Load and return config
    with open(config_path) as f:
        return yaml.safe_load(f)