hangshuo652
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
try:
|
|
import tomllib
|
|
except ImportError:
|
|
try:
|
|
import tomli as tomllib
|
|
except ImportError:
|
|
tomllib = None
|
|
|
|
|
|
@dataclass
|
|
class Config:
|
|
project_name: str = ""
|
|
copybook_paths: list[str] = field(default_factory=lambda: ["./copybooks"])
|
|
dialect: str = "ibm"
|
|
llm_model: str = "gpt-4o-mini"
|
|
llm_timeout: int = 15
|
|
llm_cache_dir: str = ".cache/llm"
|
|
coverage_default: str = "boundary"
|
|
rounding_mode: str = "TRUNCATE"
|
|
tolerance: float = 0.01
|
|
runner_mode: str = "native"
|
|
spark_master: str = "local[*]"
|
|
spark_input_format: str = "json"
|
|
num_records: int = 1000
|
|
branch_pass: float = 0.80
|
|
max_llm_cost: float = 0.50
|
|
|
|
@classmethod
|
|
def from_toml(cls, path: str = "aurak.toml") -> "Config":
|
|
if tomllib is None:
|
|
return cls()
|
|
with open(path, "rb") as f:
|
|
data = tomllib.load(f)
|
|
c = cls()
|
|
p = data.get("project", {})
|
|
c.project_name = p.get("name", c.project_name)
|
|
c.copybook_paths = p.get("copybook_paths", c.copybook_paths)
|
|
c.dialect = p.get("dialect", c.dialect)
|
|
llm = data.get("llm", {})
|
|
c.llm_model = llm.get("model", c.llm_model)
|
|
c.llm_timeout = llm.get("timeout", c.llm_timeout)
|
|
c.llm_cache_dir = llm.get("cache_dir", c.llm_cache_dir)
|
|
cov = data.get("coverage", {})
|
|
c.coverage_default = cov.get("default_target", c.coverage_default)
|
|
cmp = data.get("comparison", {})
|
|
c.rounding_mode = cmp.get("rounding_mode", c.rounding_mode)
|
|
c.tolerance = cmp.get("default_tolerance", c.tolerance)
|
|
r = data.get("runner", {})
|
|
c.runner_mode = r.get("mode", c.runner_mode)
|
|
sp = data.get("spark", {})
|
|
c.spark_master = sp.get("master", c.spark_master)
|
|
c.spark_input_format = sp.get("input_format", c.spark_input_format)
|
|
c.num_records = sp.get("num_records", c.num_records)
|
|
return c
|