C / core /config_loader.py
linxinhua's picture
key-pool rotation: per-session OpenAI key assignment from OPENAI_KEY_01..10 (core/config_loader.py)
9d5f105 verified
Raw
History Blame Contribute Delete
3.51 kB
"""Load config.yaml and initialize shared settings (API clients, OS detection, paths)."""
import os
import platform
import threading
import yaml
from openai import OpenAI
from dotenv import load_dotenv
from huggingface_hub import HfApi
from pathlib import Path
BASE_DIR = Path(__file__).resolve().parent.parent # project root (one level up from core/)
with open(BASE_DIR / "config.yaml") as f:
CONFIG = yaml.safe_load(f)
MODEL = CONFIG["model"]
PATHS = CONFIG["paths"]
REASONING_EFFORT = CONFIG.get("reasoning_effort", "medium")
VERBOSITY = CONFIG.get("verbosity", "medium")
# Load API key: .env for local dev, HF Secrets for Spaces
env_path = BASE_DIR / PATHS["env_file"]
if env_path.exists():
load_dotenv(env_path)
# OpenAI client pool for round-robin per-session key assignment.
# Loads OPENAI_KEY_01 .. OPENAI_KEY_10 from environment. Falls back to API_KEY
# (single client, no rotation) if no pool keys are present.
_pool_pairs = []
for _i in range(1, 11):
_k = os.getenv(f"OPENAI_KEY_{_i:02d}")
if _k:
_pool_pairs.append((_i, _k))
KEY_POOL = []
if _pool_pairs:
KEY_POOL = [(idx, OpenAI(api_key=k, max_retries=5)) for idx, k in _pool_pairs]
print(f"[key_pool] loaded {len(KEY_POOL)} OPENAI_KEY_NN secrets for per-session rotation", flush=True)
else:
_fallback = os.getenv("API_KEY")
if _fallback:
KEY_POOL = [(0, OpenAI(api_key=_fallback, max_retries=5))]
print("[key_pool] no OPENAI_KEY_NN secrets found; using single API_KEY (no rotation)", flush=True)
else:
print("[key_pool] WARNING: no OpenAI key configured", flush=True)
_counter = 0
_counter_lock = threading.Lock()
def get_next_client():
"""Round-robin OpenAI client picker. Returns (key_idx, client)."""
global _counter
if not KEY_POOL:
raise RuntimeError("no OpenAI client configured (set OPENAI_KEY_01..10 or API_KEY)")
with _counter_lock:
idx_in_pool = _counter % len(KEY_POOL)
_counter += 1
return KEY_POOL[idx_in_pool]
client = KEY_POOL[0][1] if KEY_POOL else None
# HuggingFace logging setup. RUN_MODE env var picks the bucket
# (test|pilot|prod) appended to mode_prefix to form the final HF dataset path.
LOG_CONFIG = CONFIG.get("logging", {})
HF_TOKEN = os.getenv(LOG_CONFIG.get("hf_token_env", "HF_access"))
HF_DATASET = LOG_CONFIG.get("hf_dataset")
hf_api = HfApi(token=HF_TOKEN) if HF_TOKEN and HF_DATASET else None
RUN_MODE = os.getenv("RUN_MODE", "test")
LOG_MODE_PREFIX = LOG_CONFIG.get("mode_prefix", "HOT/socratic")
LOG_PATH = f"{LOG_MODE_PREFIX}/{RUN_MODE}"
CONDITION = LOG_CONFIG.get("condition", "socratic")
SCHEMA_VERSION = LOG_CONFIG.get("schema_version", "1.0")
CASE_ID = LOG_CONFIG.get("case_id", "unknown")
# OS detection
IS_WINDOWS = platform.system() == "Windows"
# Load system prompt
_prompt_path = BASE_DIR / PATHS.get("system_prompt", "prompts/system_prompt.md")
SYSTEM_PROMPT = _prompt_path.read_text(encoding="utf-8") if _prompt_path.exists() else ""
def _load_prompt(key, default_relpath):
"""Read a prompt markdown file. Returns "" if the file is missing so the
app still boots (with reduced behaviour) rather than crashing."""
relpath = PATHS.get(key, default_relpath)
path = BASE_DIR / relpath
if not path.exists():
print(f"[config_loader] Warning: prompt file missing at {path}", flush=True)
return ""
return path.read_text(encoding="utf-8")
PERFECT_ANSWER_PROMPT = _load_prompt("perfect_answer_prompt", "prompts/perfect_answer_prompt.md")