TeamGenKI/Inference-API

1"""Utility functions for the inference API."""2import json3import logging4import os5import re6from pathlib import Path7from typing import Dict, Any8 9import yaml10 11 12def extract_json(text: str) -> Dict[str, Any]:13    """Extract JSON from text that might contain other content.14 15    Handles cases like:16    - Clean JSON: {"key": "value"}17    - JSON with prefix: Sure! Here's your JSON: {"key": "value"}18    - JSON with suffix: {"key": "value"} Let me know if you need anything else!19    """20    # Find anything that looks like a JSON object21    json_pattern = r'\{(?:[^{}]|(?R))*\}'22    matches = re.finditer(json_pattern, text)23 24    # Try each match until we find valid JSON25    for match in matches:26        try:27            potential_json = match.group()28            parsed = json.loads(potential_json)29            return parsed30        except json.JSONDecodeError:31            continue32 33    # If we couldn't find any valid JSON, raise an error34    raise ValueError("No valid JSON found in response")35 36def load_config():37    """38    Load configuration from config files in the resources directory.39    Uses CONFIG_ENV environment variable to determine which config to load.40    Defaults to 'local' if no environment is specified.41    """42    # Get environment name from env var, default to 'local'43    env_name = os.environ.get("CONFIG_ENV", "local")44 45    # Construct path to resources directory and config file46    resources_dir = Path(__file__).parent / "resources"47    config_path = resources_dir / f"{env_name}_config.yaml"48 49    # Create resources directory if it doesn't exist50    resources_dir.mkdir(exist_ok=True)51 52    # Check if config file exists53    if not config_path.exists():54        logging.warning(f"Config file {config_path} not found, falling back to local_config.yaml")55        config_path = resources_dir / "local_config.yaml"56 57        # If even local config doesn't exist, raise error58        if not config_path.exists():59            raise FileNotFoundError(60                f"No configuration file found at {config_path}. "61                "Please ensure at least local_config.yaml exists in the resources directory."62            )63 64    # Load and return config65    with open(config_path) as f:66        return yaml.safe_load(f)