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