|
22 | 22 | from pathlib import Path |
23 | 23 | from typing import Iterable, List, Optional |
24 | 24 |
|
25 | | -import yaml |
| 25 | +from ruamel.yaml import YAML |
| 26 | +from ruamel.yaml.error import YAMLError |
26 | 27 |
|
27 | 28 | from agentops.agent.config import SpecConformanceCheckConfig |
28 | 29 | from agentops.agent.findings import Category, Finding, Severity |
@@ -297,8 +298,8 @@ def _check_agent_drift(workspace: Path, doc: SpecDocument) -> List[Finding]: |
297 | 298 | if not run_yaml.exists(): |
298 | 299 | return [] |
299 | 300 | try: |
300 | | - raw = yaml.safe_load(run_yaml.read_text(encoding="utf-8")) |
301 | | - except (OSError, yaml.YAMLError): |
| 301 | + raw = YAML(typ="safe").load(run_yaml.read_text(encoding="utf-8")) |
| 302 | + except (OSError, YAMLError): |
302 | 303 | return [] |
303 | 304 | if not isinstance(raw, dict): |
304 | 305 | return [] |
@@ -342,8 +343,8 @@ def _collect_evaluator_names(workspace: Path) -> set[str]: |
342 | 343 | return out |
343 | 344 | for p in bundles_dir.glob("*.y*ml"): |
344 | 345 | try: |
345 | | - raw = yaml.safe_load(p.read_text(encoding="utf-8")) |
346 | | - except (OSError, yaml.YAMLError): |
| 346 | + raw = YAML(typ="safe").load(p.read_text(encoding="utf-8")) |
| 347 | + except (OSError, YAMLError): |
347 | 348 | continue |
348 | 349 | if not isinstance(raw, dict): |
349 | 350 | continue |
|
0 commit comments