1
0
Fork 0
taleweave-ai/taleweave/schema.py

87 lines
1.9 KiB
Python

from argparse import ArgumentParser
from json import dumps
from pydantic import TypeAdapter
from taleweave.models.entity import Character, Item, Portal, Room, World
from taleweave.utils.file import load_yaml
from taleweave.utils.world import describe_entity
MODELS = {
"character": Character,
"item": Item,
"portal": Portal,
"room": Room,
"world": World,
}
def parse_args():
parser = ArgumentParser(
description="Generate or validate a schema for the TaleWeave AI models"
)
subparsers = parser.add_subparsers(dest="command", required=True)
subparsers.required = True
generate_parser = subparsers.add_parser(
"generate", help="Generate a schema for a model"
)
generate_parser.add_argument(
"model",
type=str,
choices=list(MODELS.keys()),
help="The name of the model to generate a schema for",
)
validate_parser = subparsers.add_parser(
"validate", help="Validate a JSON file against a model schema"
)
validate_parser.add_argument(
"model",
type=str,
choices=list(MODELS.keys()),
help="The name of the model to validate against",
)
validate_parser.add_argument(
"file",
type=str,
help="The path to the JSON file to validate",
)
return parser.parse_args()
def command_generate(args):
model = MODELS[args.model]
schema = TypeAdapter(model).json_schema()
print(dumps(schema, indent=2))
def command_validate(args):
model = MODELS[args.model]
with open(args.file, "r") as file:
data = load_yaml(file)
entity = model(**data)
if isinstance(entity, World):
print(entity)
else:
print(describe_entity(entity))
COMMANDS = {
"generate": command_generate,
"validate": command_validate,
}
def main():
args = parse_args()
command = COMMANDS[args.command]
command(args)
if __name__ == "__main__":
main()