wuenlp_tools.config

  1from __future__ import annotations
  2
  3import json
  4import os
  5from dataclasses import dataclass
  6from pathlib import Path
  7from typing import Any, Dict, Optional, Union
  8
  9from wuenlp_tools.utils import _UserPath
 10
 11__all__ = [
 12    "APIConfig",
 13    "api_config",
 14    "set_api_base_url",
 15    "get_api_base_url",
 16    "get_endpoint_url",
 17    "LightLLMProxyConfig",
 18    "LocalLLMConfig",
 19    "WUENLP_DATA_DIR",
 20    "read_llm_config",
 21    "write_llm_config",
 22    "get_lightllm_config",
 23    "get_local_llm_config",
 24    "get_pull_lock_dir",
 25    "get_litellm_metadata",
 26    "configure_lightllm_proxy",
 27    "show_lightllm_status",
 28    "get_ssc_repo_root",
 29    "get_ssc_scene_model_path",
 30    "get_ssc_suspense_model_root",
 31    "read_ssc_config",
 32    "write_ssc_config",
 33]
 34
 35
 36class APIConfig:
 37    """Global configuration for WueNLP API endpoints."""
 38
 39    _instance: Optional["APIConfig"] = None
 40
 41    def __init__(self):
 42        self._base_url: str = os.getenv(
 43            "WUENLP_API_BASE_URL", "https://wuenlp-api.professor-x.de"
 44        )
 45
 46    @classmethod
 47    def get_instance(cls) -> "APIConfig":
 48        """Get singleton instance of APIConfig."""
 49        if cls._instance is None:
 50            cls._instance = cls()
 51        return cls._instance
 52
 53    @property
 54    def base_url(self) -> str:
 55        """Get the base URL for WueNLP API."""
 56        return self._base_url
 57
 58    @base_url.setter
 59    def base_url(self, value: str) -> None:
 60        """Set the base URL for WueNLP API."""
 61        if value.endswith("/"):
 62            value = value.rstrip("/")
 63        self._base_url = value
 64
 65    def get_endpoint_url(self, endpoint: str) -> str:
 66        """Get full URL for a specific endpoint."""
 67        if not endpoint.startswith("/"):
 68            endpoint = "/" + endpoint
 69        if not endpoint.endswith("/"):
 70            endpoint = endpoint + "/"
 71        return f"{self._base_url}{endpoint}"
 72
 73
 74# Global instance
 75api_config = APIConfig.get_instance()
 76
 77
 78def set_api_base_url(url: str) -> None:
 79    """Set the global API base URL."""
 80    api_config.base_url = url
 81
 82
 83def get_api_base_url() -> str:
 84    """Get the global API base URL."""
 85    return api_config.base_url
 86
 87
 88def get_endpoint_url(endpoint: str) -> str:
 89    """Get full URL for a specific endpoint."""
 90    return api_config.get_endpoint_url(endpoint)
 91
 92
 93# LLM Configuration
 94WUENLP_DATA_DIR = _UserPath.home() / ".local/share/wuenlp-tools"
 95LLM_CONFIG_PATH = WUENLP_DATA_DIR / "llm_config.json"
 96DEFAULT_PULL_LOCK_DIR = WUENLP_DATA_DIR / "pull-locks"
 97PROJECT_LLM_CONFIG_CANDIDATES = (
 98    ".wuenlp_llm_config.json",
 99    "wuenlp_llm_config.json",
100)
101
102
103@dataclass
104class LightLLMProxyConfig:
105    enabled: bool = False
106    base_url: str = "http://localhost:8080/v1/"
107    model_prefixes: Optional[Dict[str, str]] = None
108
109    def __post_init__(self):
110        if self.model_prefixes is None:
111            self.model_prefixes = {
112                "openai": "openai/",
113                "openrouter": "openrouter/",
114                "ollama": "ollama/",
115                "local": "ollama/",
116            }
117
118
119@dataclass
120class LocalLLMConfig:
121    base_url: Optional[str] = None
122    port: Union[str, int] = "auto"
123    auto_pull: bool = True
124    manage_process: Union[str, bool] = "auto"
125    pull_lock_dir: Optional[str] = None
126
127
128def read_llm_config() -> dict:
129    """Read LLM config from JSON file. Returns empty dict if file doesn't exist."""
130    if os.path.isfile(LLM_CONFIG_PATH):
131        try:
132            with open(LLM_CONFIG_PATH, "r") as f:
133                return json.load(f)
134        except json.JSONDecodeError:
135            return {}
136    return {}
137
138
139def write_llm_config(config: dict):
140    """Write the LLM config dictionary to the JSON file."""
141    os.makedirs(os.path.dirname(LLM_CONFIG_PATH), exist_ok=True)
142    with open(LLM_CONFIG_PATH, "w") as f:
143        json.dump(config, f, indent=2)
144
145
146def get_lightllm_config() -> LightLLMProxyConfig:
147    """Get LightLLM proxy configuration."""
148    config = read_llm_config()
149    proxy_config = config.get("lightllm_proxy", {})
150    return LightLLMProxyConfig(**proxy_config)
151
152
153def get_local_llm_config() -> LocalLLMConfig:
154    """Get local Ollama configuration."""
155    config = read_llm_config()
156    local_config = config.get("local_llm", {})
157    return LocalLLMConfig(**local_config)
158
159
160def get_pull_lock_dir() -> Path:
161    """Directory for cross-job Ollama pull locks."""
162    env_dir = os.environ.get("WUENLP_PULL_LOCK_DIR")
163    if env_dir:
164        return Path(os.path.expanduser(env_dir))
165
166    cfg = get_local_llm_config()
167    if cfg.pull_lock_dir:
168        return Path(os.path.expanduser(cfg.pull_lock_dir))
169
170    ollama_models = os.environ.get("OLLAMA_MODELS")
171    if ollama_models and _is_slurm():
172        return Path(os.path.expanduser(ollama_models)) / ".wuenlp-pull-locks"
173
174    return DEFAULT_PULL_LOCK_DIR
175
176
177def _is_slurm() -> bool:
178    return bool(os.environ.get("SLURM_JOB_ID"))
179
180
181def _read_project_llm_config() -> dict:
182    """
183    Read project-level LLM config from the current working directory.
184
185    The first existing file from PROJECT_LLM_CONFIG_CANDIDATES is used.
186    Returns an empty dict if no file is found or if parsing fails.
187    """
188    cwd = os.getcwd()
189    for filename in PROJECT_LLM_CONFIG_CANDIDATES:
190        candidate = os.path.join(cwd, filename)
191        if os.path.isfile(candidate):
192            try:
193                with open(candidate, "r") as f:
194                    return json.load(f)
195            except json.JSONDecodeError:
196                return {}
197    return {}
198
199
200def get_litellm_metadata() -> Dict[str, Any]:
201    """
202    Get merged LiteLLM metadata configuration.
203
204    Order of precedence (later overrides earlier):
205    1. System-level config at LLM_CONFIG_PATH, key "litellm_metadata"
206    2. Project-level config in the current working directory
207       (".wuenlp_llm_config.json" or "wuenlp_llm_config.json"),
208       key "litellm_metadata"
209
210    This is intended for generic tags/metadata that should end up in LiteLLM logs.
211    """
212    system_config = read_llm_config()
213    system_meta = system_config.get("litellm_metadata", {}) or {}
214
215    project_config = _read_project_llm_config()
216    project_meta = project_config.get("litellm_metadata", {}) or {}
217
218    merged: Dict[str, Any] = {}
219    merged.update(system_meta)
220    merged.update(project_meta)
221    return merged
222
223
224def configure_lightllm_proxy():
225    """Interactive configuration helper for LightLLM proxy settings."""
226    print("🔧 LightLLM Proxy Configuration")
227    print("=" * 40)
228
229    # Load current config
230    current_config = read_llm_config()
231    current_proxy = current_config.get("lightllm_proxy", {})
232
233    print(
234        f"\nCurrent status: {'ENABLED' if current_proxy.get('enabled', False) else 'DISABLED'}"
235    )
236    if current_proxy.get("base_url"):
237        print(f"Current proxy URL: {current_proxy['base_url']}")
238
239    # Ask if user wants to enable/disable
240    while True:
241        enable_input = input("\nEnable LightLLM proxy? [y/n]: ").strip().lower()
242        if enable_input in ["y", "yes", "true", "1"]:
243            enabled = True
244            break
245        elif enable_input in ["n", "no", "false", "0"]:
246            enabled = False
247            break
248        else:
249            print("Please enter 'y' for yes or 'n' for no")
250
251    if not enabled:
252        # Disable proxy
253        new_config = current_config.copy()
254        new_config["lightllm_proxy"] = {"enabled": False}
255        write_llm_config(new_config)
256        print("✅ LightLLM proxy disabled")
257        return
258
259    # Configure proxy settings
260    print("\n📝 Configuring proxy settings...")
261
262    # Base URL
263    current_url = current_proxy.get("base_url", "http://localhost:8080/v1/")
264    base_url = input(f"Proxy base URL [{current_url}]: ").strip()
265    if not base_url:
266        base_url = current_url
267
268    # Ensure URL ends with /v1/
269    if not base_url.endswith("/"):
270        base_url += "/"
271    if not base_url.endswith("v1/"):
272        if base_url.endswith("/"):
273            base_url += "v1/"
274        else:
275            base_url += "/v1/"
276
277    # Model prefixes
278    current_prefixes = current_proxy.get(
279        "model_prefixes",
280        {"openai": "openai/", "openrouter": "openrouter/", "ollama": "ollama/"},
281    )
282
283    print("\n🏷️  Configure model prefixes (press Enter to keep current):")
284
285    model_prefixes = {}
286    for provider in ["openai", "openrouter", "ollama"]:
287        current_prefix = current_prefixes.get(provider, f"{provider}/")
288        prefix = input(f"{provider.capitalize()} prefix [{current_prefix}]: ").strip()
289        if not prefix:
290            prefix = current_prefix
291        model_prefixes[provider] = prefix
292
293    # Build final config
294    new_proxy_config = {
295        "enabled": True,
296        "base_url": base_url,
297        "model_prefixes": model_prefixes,
298    }
299
300    new_config = current_config.copy()
301    new_config["lightllm_proxy"] = new_proxy_config
302
303    # Show summary
304    print("\n📋 Configuration Summary:")
305    print("   Status: ENABLED")
306    print(f"   Base URL: {base_url}")
307    print("   Model prefixes:")
308    for provider, prefix in model_prefixes.items():
309        print(f"     {provider}: {prefix}")
310
311    # Confirm and save
312    while True:
313        confirm = input("\nSave this configuration? [y/n]: ").strip().lower()
314        if confirm in ["y", "yes"]:
315            write_llm_config(new_config)
316            print(f"✅ Configuration saved to {LLM_CONFIG_PATH}")
317
318            # Check for API key
319            from wuenlp_tools.keys import LIGHTLLM_API_KEY
320
321            try:
322                # This will prompt for the key if not set
323                key_check = str(LIGHTLLM_API_KEY)
324                if key_check:
325                    print("✅ LightLLM API key is configured")
326                else:
327                    print("⚠️  No LightLLM API key set (may not be required)")
328            except Exception:
329                print("⚠️  No LightLLM API key set (may not be required)")
330
331            print("\n🎉 LightLLM proxy configuration complete!")
332            print(f"   All LLM requests will now be routed through: {base_url}")
333            break
334        elif confirm in ["n", "no"]:
335            print("❌ Configuration cancelled")
336            break
337        else:
338            print("Please enter 'y' to save or 'n' to cancel")
339
340
341def show_lightllm_status():
342    """Display current LightLLM proxy configuration status."""
343    config = get_lightllm_config()
344
345    print("🔍 LightLLM Proxy Status")
346    print("=" * 25)
347
348    if config.enabled:
349        print("✅ Status: ENABLED")
350        print(f"🌐 Base URL: {config.base_url}")
351        print("🏷️  Model prefixes:")
352        if config.model_prefixes:
353            for provider, prefix in config.model_prefixes.items():
354                print(f"   {provider}: {prefix}")
355
356        # Check API key
357        from wuenlp_tools.keys import LIGHTLLM_API_KEY
358
359        try:
360            key_check = str(LIGHTLLM_API_KEY)
361            key_status = "SET" if key_check else "NOT SET"
362        except Exception:
363            key_status = "NOT SET"
364        print(f"🔑 API Key: {key_status}")
365
366    else:
367        print("❌ Status: DISABLED")
368        print("💡 Run configure_lightllm_proxy() to enable")
369
370
371from wuenlp_tools.ssc_config import (  # noqa: E402
372    get_ssc_repo_root,
373    get_ssc_scene_model_path,
374    get_ssc_suspense_model_root,
375    read_ssc_config,
376    write_ssc_config,
377)