MetaGPT

1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
"""
4
@Time    : 2024/1/4 10:18
5
@Author  : alexanderwu
6
@File    : YamlModel.py
7
"""
8
from pathlib import Path
9
from typing import Dict, Optional
10

11
import yaml
12
from pydantic import BaseModel, model_validator
13

14

15
class YamlModel(BaseModel):
16
    """Base class for yaml model"""
17

18
    extra_fields: Optional[Dict[str, str]] = None
19

20
    @classmethod
21
    def read_yaml(cls, file_path: Path, encoding: str = "utf-8") -> Dict:
22
        """Read yaml file and return a dict"""
23
        if not file_path.exists():
24
            return {}
25
        with open(file_path, "r", encoding=encoding) as file:
26
            return yaml.safe_load(file)
27

28
    @classmethod
29
    def from_yaml_file(cls, file_path: Path) -> "YamlModel":
30
        """Read yaml file and return a YamlModel instance"""
31
        return cls(**cls.read_yaml(file_path))
32

33
    def to_yaml_file(self, file_path: Path, encoding: str = "utf-8") -> None:
34
        """Dump YamlModel instance to yaml file"""
35
        with open(file_path, "w", encoding=encoding) as file:
36
            yaml.dump(self.model_dump(), file)
37

38

39
class YamlModelWithoutDefault(YamlModel):
40
    """YamlModel without default values"""
41

42
    @model_validator(mode="before")
43
    @classmethod
44
    def check_not_default_config(cls, values):
45
        """Check if there is any default config in config2.yaml"""
46
        if any(["YOUR" in v for v in values]):
47
            raise ValueError("Please set your config in config2.yaml")
48
        return values
49