src/data_agent_baseline/config.py

from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path

import yaml

PROJECT_ROOT = Path(__file__).resolve().parents[2]


def _default_dataset_root() -> Path:
    return PROJECT_ROOT / "data" / "public" / "input"


def _default_run_output_dir() -> Path:
    return PROJECT_ROOT / "artifacts" / "runs"


@dataclass(frozen=True, slots=True)
class DatasetConfig:
    root_path: Path = field(default_factory=_default_dataset_root)


@dataclass(frozen=True, slots=True)
class AgentConfig:
    # OpenAI-style config
    model: str = "gpt-4.1-mini"
    api_base: str = "https://api.openai.com/v1"
    api_key: str = ""
    max_steps: int = 16 
    temperature: float = 0.0

    # Azure OpenAI-style (NIQ) config
    azure_endpoint: str = ""
    azure_api_key: str = ""
    azure_api_version: str = ""
    azure_deployment: str = ""
    azure_custom_headers: dict[str, str] = field(default_factory=dict)

@dataclass(frozen=True, slots=True)
class LoggingConfig:
    log_debug: bool = False


@dataclass(frozen=True, slots=True)
class EvaluationConfig:
    """Phase 15: Run + evaluation workflow configuration."""
    enabled: bool = False
    gold_root: Path | None = None
    run_after_execution: bool = False
    fail_if_ground_truth_missing: bool = False
    lambda_penalty: float = 0.1
    mode: str = "standard"


@dataclass(frozen=True, slots=True)
class FeatureFlagConfig:
    """Feature flags for experimental capabilities."""
    # False: use existing architecture (baseline)
    # True:  use multi-agent architecture (Adaptive Analyst Team)
    # Supports side-by-side evaluation when both values are tested.
    enable_adaptive_analyst_team: bool = True  # ENABLED for KDD Creative Track


@dataclass(frozen=True, slots=True)
class RunConfig:
    output_dir: Path = field(default_factory=_default_run_output_dir)
    run_id: str | None = None
    max_workers: int = 4
    task_timeout_seconds: int = 600


@dataclass(frozen=True, slots=True)
class AppConfig:
    dataset: DatasetConfig = field(default_factory=DatasetConfig)
    agent: AgentConfig = field(default_factory=AgentConfig)
    run: RunConfig = field(default_factory=RunConfig)
    logging: LoggingConfig = field(default_factory=LoggingConfig)
    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
    feature_flags: FeatureFlagConfig = field(default_factory=FeatureFlagConfig)


def _path_value(raw_value: str | None, default_value: Path) -> Path:
    if not raw_value:
        return default_value
    candidate = Path(raw_value)
    if candidate.is_absolute():
        return candidate
    return (PROJECT_ROOT / candidate).resolve()


def load_app_config(config_path: Path) -> AppConfig:
    payload = yaml.safe_load(config_path.read_text()) or {}
    dataset_defaults = DatasetConfig()
    agent_defaults = AgentConfig()
    run_defaults = RunConfig()
    logging_defaults = LoggingConfig()
    evaluation_defaults = EvaluationConfig()

    dataset_payload = payload.get("dataset", {})
    agent_payload = payload.get("agent", {})
    run_payload = payload.get("run", {})
    logging_payload = payload.get("logging", {})
    evaluation_payload = payload.get("evaluation", {})

    dataset_config = DatasetConfig(
        root_path=_path_value(dataset_payload.get("root_path"), dataset_defaults.root_path),
    )

    # Parse custom headers for Azure
    azure_custom_headers = {}
    if "azure_custom_headers" in agent_payload:
        azure_custom_headers = dict(agent_payload.get("azure_custom_headers", {}))
    
    # Load from environment variables if not specified in config
    azure_endpoint = str(agent_payload.get("azure_endpoint", agent_defaults.azure_endpoint))
    if not azure_endpoint:
        azure_endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT", "")
    
    azure_api_key = str(agent_payload.get("azure_api_key", agent_defaults.azure_api_key))
    if not azure_api_key:
        azure_api_key = os.environ.get("AZURE_OPENAI_API_KEY", "")
    
    azure_api_version = str(agent_payload.get("azure_api_version", agent_defaults.azure_api_version))
    if not azure_api_version:
        azure_api_version = os.environ.get("AZURE_OPENAI_API_VERSION", "")
    
    azure_deployment = str(agent_payload.get("azure_deployment", agent_defaults.azure_deployment))
    if not azure_deployment:
        azure_deployment = os.environ.get("AZURE_OPENAI_DEPLOYMENT", "")
    
    # Load custom header from environment if specified
    if not azure_custom_headers:
        niq_consumer = os.environ.get("X_NIQ_CIS_CONSUMER", "")
        if niq_consumer:
            azure_custom_headers["X-NIQ-CIS-CONSUMER"] = niq_consumer
    
    # Load OpenAI config from environment if not specified
    api_key = str(agent_payload.get("api_key", agent_defaults.api_key))
    if not api_key:
        api_key = os.environ.get("OPENAI_API_KEY", "")

    agent_config = AgentConfig(
        model=str(agent_payload.get("model", agent_defaults.model)),
        api_base=str(agent_payload.get("api_base", agent_defaults.api_base)),
        api_key=api_key,
        max_steps=int(agent_payload.get("max_steps", agent_defaults.max_steps)),
        temperature=float(agent_payload.get("temperature", agent_defaults.temperature)),
        # Azure config
        azure_endpoint=azure_endpoint,
        azure_api_key=azure_api_key,
        azure_api_version=azure_api_version,
        azure_deployment=azure_deployment,
        azure_custom_headers=azure_custom_headers,
    )
    raw_run_id = run_payload.get("run_id")
    run_id = run_defaults.run_id
    if raw_run_id is not None:
        normalized_run_id = str(raw_run_id).strip()
        run_id = normalized_run_id or None

    run_config = RunConfig(
        output_dir=_path_value(run_payload.get("output_dir"), run_defaults.output_dir),
        run_id=run_id,
        max_workers=int(run_payload.get("max_workers", run_defaults.max_workers)),
        task_timeout_seconds=int(run_payload.get("task_timeout_seconds", run_defaults.task_timeout_seconds)),
    )
    
    logging_config = LoggingConfig(
        log_debug=bool(logging_payload.get("log_debug", logging_defaults.log_debug)),
    )

    # Parse evaluation config
    eval_mode = str(evaluation_payload.get("mode", evaluation_defaults.mode))
    if eval_mode not in {"standard", "verbose", "research"}:
        raise ValueError(
            f"Invalid evaluation mode: {eval_mode!r}. "
            "Must be one of: standard, verbose, research"
        )
    
    eval_gold_root = evaluation_payload.get("gold_root")
    eval_gold_root_path: Path | None = None
    if eval_gold_root is not None:
        eval_gold_root_path = _path_value(str(eval_gold_root), Path("."))
    
    evaluation_config = EvaluationConfig(
        enabled=bool(evaluation_payload.get("enabled", evaluation_defaults.enabled)),
        gold_root=eval_gold_root_path,
        run_after_execution=bool(evaluation_payload.get("run_after_execution", evaluation_defaults.run_after_execution)),
        fail_if_ground_truth_missing=bool(evaluation_payload.get("fail_if_ground_truth_missing", evaluation_defaults.fail_if_ground_truth_missing)),
        lambda_penalty=float(evaluation_payload.get("lambda_penalty", evaluation_defaults.lambda_penalty)),
        mode=eval_mode,
    )

    feature_flags_payload = payload.get("feature_flags", {})
    feature_flag_defaults = FeatureFlagConfig()
    # Also check environment variable override
    aat_env = os.environ.get("ENABLE_ADAPTIVE_ANALYST_TEAM", "").lower()
    enable_aat = bool(feature_flags_payload.get(
        "enable_adaptive_analyst_team",
        feature_flag_defaults.enable_adaptive_analyst_team,
    ))
    if aat_env in ("1", "true", "yes"):
        enable_aat = True
    elif aat_env in ("0", "false", "no"):
        enable_aat = False
    feature_flags_config = FeatureFlagConfig(enable_adaptive_analyst_team=enable_aat)

    return AppConfig(
        dataset=dataset_config,
        agent=agent_config,
        run=run_config,
        logging=logging_config,
        evaluation=evaluation_config,
        feature_flags=feature_flags_config,
    )


def load_agent_config_from_env() -> AgentConfig:
    """
    Load agent configuration from environment variables (for evaluation system).
    
    This function is used during competition evaluation when the container
    receives MODEL_API_URL, MODEL_API_KEY, and MODEL_NAME from the evaluation system.
    
    Required environment variables:
    - MODEL_API_URL: OpenAI-compatible API endpoint
    - MODEL_API_KEY: API key for authentication
    - MODEL_NAME: Model name (default: qwen3.5-35b-a3b)
    
    Returns:
        AgentConfig with settings from environment variables
    """
    model_api_url = os.getenv("MODEL_API_URL", "")
    model_api_key = os.getenv("MODEL_API_KEY", "")
    model_name = os.getenv("MODEL_NAME", "qwen3.5-35b-a3b")
    
    if not model_api_url:
        raise ValueError("MODEL_API_URL environment variable is required")
    if not model_api_key:
        raise ValueError("MODEL_API_KEY environment variable is required")
    
    return AgentConfig(
        model=model_name,
        api_base=model_api_url,
        api_key=model_api_key,
        max_steps=16,
        temperature=0.0,
        azure_endpoint="",
        azure_api_key="",
        azure_api_version="",
        azure_deployment="",
        azure_custom_headers={},
    )