87 lines
1.9 KiB
Python
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()
|