diff --git a/README.md b/README.md index d1ab5d9a..27374221 100644 --- a/README.md +++ b/README.md @@ -726,6 +726,7 @@ options: -v, --verbose make it verbose -d, --dump-configuration dump actual configuration into JSON file and quit + -s, --dump-schema dump configuration schema into OpenAPI-compatible file and quit -c CONFIG_FILE, --config CONFIG_FILE path to configuration file (default: lightspeed-stack.yaml) diff --git a/src/lightspeed_stack.py b/src/lightspeed_stack.py index fb6f4f8f..8df663c0 100644 --- a/src/lightspeed_stack.py +++ b/src/lightspeed_stack.py @@ -8,12 +8,14 @@ import os from argparse import ArgumentParser + from rich.logging import RichHandler from log import get_logger from configuration import configuration from runners.uvicorn import start_uvicorn from runners.quota_scheduler import start_quota_scheduler +from utils import schema_dumper FORMAT = "%(message)s" logging.basicConfig( @@ -29,6 +31,7 @@ def create_argument_parser() -> ArgumentParser: The parser includes these options: - -v / --verbose: enable verbose output - -d / --dump-configuration: dump the loaded configuration to JSON and exit + - -s / --dump-schema: dump the configuration schema to OpenAPI JSON and exit - -c / --config: path to the configuration file (default "lightspeed-stack.yaml") - -g / --generate-llama-stack-configuration: generate a Llama Stack configuration from the service configuration @@ -55,6 +58,14 @@ def create_argument_parser() -> ArgumentParser: action="store_true", default=False, ) + parser.add_argument( + "-s", + "--dump-schema", + dest="dump_schema", + help="dump configuration schema into OpenAPI-compatible file and quit", + action="store_true", + default=False, + ) parser.add_argument( "-c", "--config", @@ -74,6 +85,8 @@ def main() -> None: Parses command-line arguments, loads the configured settings, and then: - If --dump-configuration is provided, writes the active configuration to configuration.json and exits (exits with status 1 on failure). + - If --dump-schema is provided, writes the active configuration schema to + schema.json and exits (exits with status 1 on failure). - If --generate-llama-stack-configuration is provided, generates and stores the Llama Stack configuration to the specified output file and exits (exits with status 1 on failure). @@ -105,6 +118,17 @@ def main() -> None: raise SystemExit(1) from e return + # -s or --dump-schema CLI flags are used to dump configuration schema + # into a JSON file that is compatible with OpenAPI schema specification + if args.dump_schema: + try: + schema_dumper.dump_schema("schema.json") + logger.info("Configuration schema dumped to schema.json") + except Exception as e: + logger.error("Failed to dump configuration schema: %s", e) + raise SystemExit(1) from e + return + # Store config path in env so each uvicorn worker can load it # (step is needed because process context isn't shared). os.environ["LIGHTSPEED_STACK_CONFIG_PATH"] = args.config_file diff --git a/src/utils/schema_dumper.py b/src/utils/schema_dumper.py new file mode 100644 index 00000000..e3c8902d --- /dev/null +++ b/src/utils/schema_dumper.py @@ -0,0 +1,79 @@ +"""Function to dump the configuration schema into OpenAPI-compatible format.""" + +import json +from pydantic.json_schema import models_json_schema + +from models.config import Configuration + + +def recursive_update(original: dict) -> dict: + """Recursively update the schema to be 100% OpenAPI-compatible. + + Parameters: + original: The original schema dictionary to transform. + Returns: + A new dictionary with OpenAPI-compatible transformations applied. + """ + new: dict = {} + for key, value in original.items(): + # recurse into sub-dictionaries + if isinstance(value, dict): + new[key] = recursive_update(original[key]) + # optional types fixes + elif ( + key == "anyOf" + and isinstance(value, list) + and len(value) >= 2 + and "type" in value[0] + and value[1]["type"] == "null" + ): + # only the first type is correct, + # we need to ignore the second one + val = value[0]["type"] + new["type"] = val + # create new attribute + new["nullable"] = True + # exclusiveMinimum attribute handling is broken + # in Pydantic - this is simple fix + elif key == "exclusiveMinimum": + new["minimum"] = value + else: + new[key] = value + return new + + +def dump_schema(filename: str) -> None: + """Dump the configuration schema into OpenAPI-compatible JSON file. + + Parameters: + - filename: str - name of file to export the schema to + + Returns: + - None + + Raises: + IOError: If the file cannot be written. + """ + with open(filename, "w", encoding="utf-8") as fout: + # retrieve the schema + _, schemas = models_json_schema( + [(model, "validation") for model in [Configuration]], + ref_template="#/components/schemas/{model}", + ) + + # fix the schema + schemas = recursive_update(schemas) + + # add all required metadata + openapi_schema = { + "openapi": "3.0.0", + "info": { + "title": "Lightspeed Core Stack", + "version": "0.3.0", + }, + "components": { + "schemas": schemas.get("$defs", {}), + }, + "paths": {}, + } + json.dump(openapi_schema, fout, indent=4)