|
15 | 15 | from __future__ import annotations
|
16 | 16 |
|
17 | 17 | import json
|
18 |
| -from typing import Any, Dict, List, Literal, Optional, Tuple |
| 18 | +import yaml |
| 19 | +from typing import Any, Dict, List, Literal, Optional, Tuple, Union |
| 20 | +from pathlib import Path |
19 | 21 |
|
20 | 22 | from pydantic import BaseModel, ValidationError, model_validator, validate_call
|
21 | 23 | from requests.exceptions import InvalidJSONError
|
@@ -127,6 +129,91 @@ def check_schema(cls, data: Dict[str, Any]) -> Dict[str, Any]:
|
127 | 129 |
|
128 | 130 | return data
|
129 | 131 |
|
| 132 | + def store_as_json(self, file_path: str) -> None: |
| 133 | + """ |
| 134 | + Save the schema configuration to a JSON file. |
| 135 | +
|
| 136 | + Args: |
| 137 | + file_path (str): The path where the schema configuration will be saved. |
| 138 | + """ |
| 139 | + with open(file_path, 'w') as f: |
| 140 | + json.dump(self.model_dump(), f, indent=2) |
| 141 | + |
| 142 | + def store_as_yaml(self, file_path: str) -> None: |
| 143 | + """ |
| 144 | + Save the schema configuration to a YAML file. |
| 145 | +
|
| 146 | + Args: |
| 147 | + file_path (str): The path where the schema configuration will be saved. |
| 148 | + """ |
| 149 | + with open(file_path, 'w') as f: |
| 150 | + yaml.dump(self.model_dump(), f, default_flow_style=False, sort_keys=False) |
| 151 | + |
| 152 | + @classmethod |
| 153 | + def from_file(cls, file_path: Union[str, Path]) -> Self: |
| 154 | + """ |
| 155 | + Load a schema configuration from a file (either JSON or YAML). |
| 156 | + |
| 157 | + The file format is automatically detected based on the file extension. |
| 158 | + |
| 159 | + Args: |
| 160 | + file_path (Union[str, Path]): The path to the schema configuration file. |
| 161 | + |
| 162 | + Returns: |
| 163 | + SchemaConfig: The loaded schema configuration. |
| 164 | + """ |
| 165 | + file_path = Path(file_path) |
| 166 | + |
| 167 | + if not file_path.exists(): |
| 168 | + raise FileNotFoundError(f"Schema file not found: {file_path}") |
| 169 | + |
| 170 | + if file_path.suffix.lower() in ['.json']: |
| 171 | + return cls.from_json(file_path) |
| 172 | + elif file_path.suffix.lower() in ['.yaml', '.yml']: |
| 173 | + return cls.from_yaml(file_path) |
| 174 | + else: |
| 175 | + raise ValueError(f"Unsupported file format: {file_path.suffix}. Use .json, .yaml, or .yml") |
| 176 | + |
| 177 | + @classmethod |
| 178 | + def from_json(cls, file_path: Union[str, Path]) -> Self: |
| 179 | + """ |
| 180 | + Load a schema configuration from a JSON file. |
| 181 | + |
| 182 | + Args: |
| 183 | + file_path (Union[str, Path]): The path to the JSON schema configuration file. |
| 184 | + |
| 185 | + Returns: |
| 186 | + SchemaConfig: The loaded schema configuration. |
| 187 | + """ |
| 188 | + with open(file_path, 'r') as f: |
| 189 | + try: |
| 190 | + data = json.load(f) |
| 191 | + return cls.model_validate(data) |
| 192 | + except json.JSONDecodeError as e: |
| 193 | + raise ValueError(f"Invalid JSON file: {e}") |
| 194 | + except ValidationError as e: |
| 195 | + raise SchemaValidationError(f"Schema validation failed: {e}") |
| 196 | + |
| 197 | + @classmethod |
| 198 | + def from_yaml(cls, file_path: Union[str, Path]) -> Self: |
| 199 | + """ |
| 200 | + Load a schema configuration from a YAML file. |
| 201 | + |
| 202 | + Args: |
| 203 | + file_path (Union[str, Path]): The path to the YAML schema configuration file. |
| 204 | + |
| 205 | + Returns: |
| 206 | + SchemaConfig: The loaded schema configuration. |
| 207 | + """ |
| 208 | + with open(file_path, 'r') as f: |
| 209 | + try: |
| 210 | + data = yaml.safe_load(f) |
| 211 | + return cls.model_validate(data) |
| 212 | + except yaml.YAMLError as e: |
| 213 | + raise ValueError(f"Invalid YAML file: {e}") |
| 214 | + except ValidationError as e: |
| 215 | + raise SchemaValidationError(f"Schema validation failed: {e}") |
| 216 | + |
130 | 217 |
|
131 | 218 | class SchemaBuilder(Component):
|
132 | 219 | """
|
@@ -281,7 +368,7 @@ async def run(self, text: str, **kwargs: Any) -> SchemaConfig:
|
281 | 368 | try:
|
282 | 369 | extracted_schema: Dict[str, Any] = json.loads(content)
|
283 | 370 | except json.JSONDecodeError as exc:
|
284 |
| - raise InvalidJSONError( |
| 371 | + raise ValueError( |
285 | 372 | "LLM response is not valid JSON."
|
286 | 373 | ) from exc
|
287 | 374 |
|
|
0 commit comments