Skip to content

Commit 49452d4

Browse files
Save/Read inferred schema
1 parent 2b14541 commit 49452d4

File tree

1 file changed

+89
-2
lines changed
  • src/neo4j_graphrag/experimental/components

1 file changed

+89
-2
lines changed

src/neo4j_graphrag/experimental/components/schema.py

Lines changed: 89 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@
1515
from __future__ import annotations
1616

1717
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
1921

2022
from pydantic import BaseModel, ValidationError, model_validator, validate_call
2123
from requests.exceptions import InvalidJSONError
@@ -127,6 +129,91 @@ def check_schema(cls, data: Dict[str, Any]) -> Dict[str, Any]:
127129

128130
return data
129131

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+
130217

131218
class SchemaBuilder(Component):
132219
"""
@@ -281,7 +368,7 @@ async def run(self, text: str, **kwargs: Any) -> SchemaConfig:
281368
try:
282369
extracted_schema: Dict[str, Any] = json.loads(content)
283370
except json.JSONDecodeError as exc:
284-
raise InvalidJSONError(
371+
raise ValueError(
285372
"LLM response is not valid JSON."
286373
) from exc
287374

0 commit comments

Comments
 (0)