diff --git a/mlx/cos_backend.py b/mlx/cos_backend.py index e2e83c1..0345412 100644 --- a/mlx/cos_backend.py +++ b/mlx/cos_backend.py @@ -2,7 +2,7 @@ import boto3 import jsonschema from jsonschema import validate - +import os class COSKVStore: def __init__(self, bucket_name, schema, cos_client=None): @@ -26,7 +26,6 @@ def put(self, key, value): validate(instance=value, schema=self.schema) except jsonschema.exceptions.ValidationError as e: raise ValueError(f"Validation error: {e.message}") - self.cos_client.put_object( Bucket=self.bucket_name, Key=key, @@ -60,19 +59,53 @@ def list_keys(self): response = self.cos_client.list_objects_v2(Bucket=self.bucket_name) return [obj['Key'] for obj in response.get('Contents', [])] +def load_schemas(schema_folder): + """ + Loads all JSON schemas from the given folder. + :param schema_folder: Path to the folder containing JSON schema files. + :return: A dictionary of schema names and their corresponding JSON objects. + """ + schemas = {} + for filename in os.listdir(schema_folder): + if filename.endswith(".json"): + filepath = os.path.join(schema_folder, filename) + with open(filepath, 'r') as f: + schemas[filename[:-5]] = json.load(f) + return schemas # Example Usage if __name__ == "__main__": - example_schema = { - "type": "object", - "properties": { - "name": {"type": "string"}, - "age": {"type": "integer"} - }, - "required": ["name", "age"] - } + schema_folder = "../schemas" + schemas = load_schemas(schema_folder) + + if "example_schema" in schemas: + example_schema = schemas["example_schema"] + cos_store = COSKVStore("my-cos-bucket", example_schema) + + test_data = { + "id": "model_123", + "name": "MyModel", + "framework": "TensorFlow", + "version": "2.10", + "description": "A simple neural network model.", + "metrics": { + "accuracy": 0.95, + "loss": 0.1 + } + } + cos_store.put("model_123", test_data) + print(cos_store.get("model_123")) + + #Example data without the metric key, which is not required + test_data_no_metrics = { + "id": "model_456", + "name": "MyOtherModel", + "framework": "PyTorch", + "version": "1.12", + "description": "Another neural network model.", + } + cos_store.put("model_456", test_data_no_metrics) + print(cos_store.get("model_456")) - cos_store = COSKVStore("my-cos-bucket", example_schema) - test_data = {"name": "John Doe", "age": 30} - cos_store.put("user_123", test_data) - print(cos_store.get("user_123")) + else: + print("Schema 'example_schema' not found.") \ No newline at end of file