Skip to content

Commit

Permalink
feat(py): generate schemas.py from genkit-schema.json using datam…
Browse files Browse the repository at this point in the history
…odel-codegen #1807 (#1808)

ISSUE: #1807

CHANGELOG:
- [ ] Add a `bin/generate_schema_types` script that generates the Pydantic
  `schemas.py` module.
- [ ] Not updating the pre-commit hooks since the generated file
  requires manual patching.
- [ ] Remove timestamp to ensure we do not treat a file with identical
content differently preventing the hassles of updating this file per
commit.
  • Loading branch information
yesudeep authored Feb 11, 2025
1 parent 889e309 commit 3a3f69f
Show file tree
Hide file tree
Showing 10 changed files with 566 additions and 216 deletions.
3 changes: 3 additions & 0 deletions py/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
# Copyright 2025 Google LLC
# SPDX-License-Identifier: Apache-2.0

35 changes: 35 additions & 0 deletions py/bin/generate_schema_types
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
#!/usr/bin/env bash
#
# Copyright 2025 Google LLC
# SPDX-License-Identifier: Apache-2.0

set -euo pipefail

TOP_DIR=$(git rev-parse --show-toplevel)
SCHEMA_FILE="$TOP_DIR/py/packages/genkit/src/genkit/core/schemas.py"

# Generate types using configuration from pyproject.toml
uv run --directory "$TOP_DIR/py" datamodel-codegen

# This isn't causing runtime errors at the moment so letting it be.
#sed -i '' '/^class Model(RootModel\[Any\]):$/,/^ root: Any$/d' "$SCHEMA_FILE"

# 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 check --fix "$SCHEMA_FILE"
uv run --directory "$TOP_DIR/py" \
ruff format "$SCHEMA_FILE"
85 changes: 85 additions & 0 deletions py/bin/sanitize_schemas.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
#!/usr/bin/env python3
#
# Copyright 2025 Google LLC
# SPDX-License-Identifier: Apache-2.0

import ast
import sys


class ModelConfigRemover(ast.NodeTransformer):
def __init__(self) -> None:
self.modified = False

def is_rootmodel_class(self, node: ast.ClassDef) -> bool:
"""Check if the class inherits from RootModel."""
for base in node.bases:
if isinstance(base, ast.Name) and base.id == 'RootModel':
return True
elif isinstance(base, ast.Subscript):
if (
isinstance(base.value, ast.Name)
and base.value.id == 'RootModel'
):
return True
return False

def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
"""Visit class definitions and remove model_config if class inherits from RootModel."""
if self.is_rootmodel_class(node):
# Filter out model_config assignments
new_body = []
for item in node.body:
if isinstance(item, ast.Assign):
targets = item.targets
if len(targets) == 1 and isinstance(targets[0], ast.Name):
if targets[0].id != 'model_config':
new_body.append(item)
else:
new_body.append(item)

if len(new_body) != len(node.body):
self.modified = True

node.body = new_body

return node


def process_file(filename: str) -> None:
"""Process a Python file to remove model_config from RootModel classes."""
with open(filename, 'r') as f:
source = f.read()

tree = ast.parse(source)

transformer = ModelConfigRemover()
modified_tree = transformer.visit(tree)

if transformer.modified:
ast.fix_missing_locations(modified_tree)
modified_source = ast.unparse(modified_tree)
with open(filename, 'w') as f:
f.write(modified_source)
print(
f'Modified {filename}: Removed model_config from RootModel classes'
)
else:
print(f'No modifications needed in {filename}')


def main() -> None:
if len(sys.argv) != 2:
print('Usage: python script.py <filename>')
sys.exit(1)

filename = sys.argv[1]
try:
process_file(filename)
except Exception as e:
print(f'Error processing {filename}: {str(e)}')
sys.exit(1)


if __name__ == '__main__':
main()
6 changes: 6 additions & 0 deletions py/captainhook.json
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,9 @@
{
"run": "pnpm i --frozen-lockfile"
},
{
"run": "py/bin/generate_schema_types"
},
{
"run": "py/bin/fmt"
},
Expand Down Expand Up @@ -81,6 +84,9 @@
{
"run": "pnpm i --frozen-lockfile"
},
{
"run": "py/bin/generate_schema_types"
},
{
"run": "py/bin/fmt"
},
Expand Down
Loading

0 comments on commit 3a3f69f

Please sign in to comment.