import only adds by default

import_schemas.py

@@ -2,6 +2,7 @@
 import replicate
 import json
 import os
+import argparse
 
 
 def format_json_file(file_path):
@@ -25,13 +26,29 @@ def format_json_files_in_directory(directory):
             format_json_file(file_path)
 
 
-with open("supported_models.json", "r", encoding="utf-8") as f:
-    supported_models = json.load(f)
+def update_schemas(update=False):
+    with open("supported_models.json", "r", encoding="utf-8") as f:
+        supported_models = json.load(f)
 
-for model in supported_models["models"]:
-    m = replicate.models.get(model)
-    with open(f"schemas/{model.replace('/', '_')}.json", "w", encoding="utf-8") as f:
-        f.write(m.json())
+    schemas_directory = "schemas"
+    existing_schemas = set(os.listdir(schemas_directory))
 
-schemas_directory = "schemas"
-format_json_files_in_directory(schemas_directory)
+    for model in supported_models["models"]:
+        schema_filename = f"{model.replace('/', '_')}.json"
+        schema_path = os.path.join(schemas_directory, schema_filename)
+
+        if update or schema_filename not in existing_schemas:
+            m = replicate.models.get(model)
+            with open(schema_path, "w", encoding="utf-8") as f:
+                f.write(m.json())
+            print(f"{'Updated' if update else 'Added'} schema for {model}")
+
+    format_json_files_in_directory(schemas_directory)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Update model schemas")
+    parser.add_argument("--update", action="store_true", help="Update all schemas, not just new ones")
+    args = parser.parse_args()
+
+    update_schemas(update=args.update)