diff --git a/config.py b/pipelines/config.py similarity index 61% rename from config.py rename to pipelines/config.py index b70df9b..4bdaadb 100644 --- a/config.py +++ b/pipelines/config.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from os import getenv from pathlib import Path import tomllib @@ -68,15 +69,54 @@ class BenchmarkConfig: return cls(**raw) +@dataclass +class OpenAIConfig: + """OpenAI API configuration.""" + + api_key: str + openai_project_id: str + openai_chat_completions_url: str + model: str + timeout_seconds: int + + @classmethod + def from_toml(cls, config_path: Path) -> OpenAIConfig: + """Load OpenAI config from a TOML file.""" + raw = tomllib.loads(config_path.read_text()).get("openai", {}) + api_key = getenv("CLOSEDAI_TOKEN") + if not api_key: + message = "CLOSEDAI_TOKEN is required" + raise KeyError(message) + return cls( + api_key=api_key, + openai_project_id=raw.get( + "openai_project_id", "proj_fQBPEXFgnS87Fk6wZwploFwE" + ), + openai_chat_completions_url=raw.get( + "openai_chat_completions_url", + "https://api.openai.com/v1/chat/completions", + ), + model=raw.get("model", "gpt-5.4-mini"), + timeout_seconds=raw.get("timeout_seconds", 60), + ) + + def get_config_dir() -> Path: - """Get the path to the config file.""" - return Path(__file__).resolve().parent.parent.parent / "config" + """Get the path to the config directory.""" + return Path(__file__).resolve().parents[2] / "config" + def default_config_path() -> Path: """Get the path to the config file.""" return get_config_dir() / "config.toml" +def get_openai_config(config_path: Path | None = None) -> OpenAIConfig: + if config_path is None: + config_path = default_config_path() + return OpenAIConfig.from_toml(config_path) + + def get_finetune_config(config_path: Path | None = None) -> FinetuneConfig: if config_path is None: config_path = default_config_path()