fix: py/bin/generate_schema_types should generate the schema types co…

…rrectly now #2002

ISSUE: #2002

CHANGELOG:
- [ ] Update the `py/bin/generate_schema_types` script to add the
  headers in `sanitize_schemas.py` module directly.

py/bin/generate_schema_types

@@ -17,17 +17,6 @@ uv run --directory "${TOP_DIR}/py" datamodel-codegen
 # Sanitize the generated schema.
 python3 "${TOP_DIR}/py/bin/sanitize_schemas.py" "${SCHEMA_FILE}"
 
-# Add a generated by `generate_schema_types` comment.
-sed -i '' '1i\
-# DO NOT EDIT: Generated by `generate_schema_types` from `genkit-schemas.json`.
-' "${SCHEMA_FILE}"
-
-# Add license header.
-addlicense \
-  -c "Google LLC" \
-  -s=only \
-  "${SCHEMA_FILE}"
-
 # Checks and formatting.
 uv run --directory "${TOP_DIR}/py" \
   ruff format "${TOP_DIR}"

py/bin/sanitize_schemas.py

@@ -5,6 +5,8 @@
 
 import ast
 import sys
+from datetime import datetime
+from pathlib import Path
 
 
 class ModelConfigRemover(ast.NodeTransformer):
@@ -47,26 +49,48 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:  # noqa: N802
         return node
 
 
+def add_header(content: str) -> str:
+    """Add the generated header to the content."""
+    header = """# Copyright {year} Google LLC
+# SPDX-License-Identifier: Apache-2.0
+
+# DO NOT EDIT: Generated by `generate_schema_types` from `genkit-schemas.json`.
+
+from __future__ import annotations
+"""
+    return header.format(year=datetime.now().year) + content
+
+
 def process_file(filename: str) -> None:
-    """Process a Python file to remove model_config from RootModel classes."""
-    with open(filename) as f:
-        source = f.read()
+    """Process a Python file to remove model_config from RootModel classes
+    and add the generated header."""
+    path = Path(filename)
+    if not path.is_file():
+        raise FileNotFoundError(f'{filename} does not exist or is not a file')
 
-    tree = ast.parse(source)
+    # Read the file
+    content = path.read_text(encoding='utf-8')
 
+    # Then process the AST
+    tree = ast.parse(content)
     transformer = ModelConfigRemover()
     modified_tree = transformer.visit(tree)
 
     if transformer.modified:
+        # Write back the modified content.
         ast.fix_missing_locations(modified_tree)
         modified_source = ast.unparse(modified_tree)
-        with open(filename, 'w') as f:
-            f.write(modified_source)
+        src = add_header(modified_source)
+        print(src)
+        # print('# foo\n' + src)
+        path.write_text(src, encoding='utf-8')
         print(
             f'Modified {filename}: Removed model_config from RootModel classes'
         )
     else:
-        print(f'No modifications needed in {filename}')
+        # Even if no AST modifications, still write back to add the header.
+        path.write_text(add_header(content), encoding='utf-8')
+        print(f'Added header to {filename}')
 
 
 def main() -> None:

py/packages/genkit/src/genkit/core/schemas.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # DO NOT EDIT: Generated by `generate_schema_types` from `genkit-schemas.json`.
+
 from __future__ import annotations
 
 from enum import Enum