Merge pull request lightspeed-core#2109 from tisnik/lcore-2755-dumper-refactoring

LCORE-2755: schema and models dumper refactoring

Author: tisnik

src/lightspeed_stack.py

@@ -14,7 +14,7 @@
 from log import get_logger, setup_logging
 from runners.quota_scheduler import start_quota_scheduler
 from runners.uvicorn import start_uvicorn
-from utils import models_dumper, schema_dumper
+from utils import config_dumper, models_dumper
 
 setup_logging()
 logger = get_logger(__name__)
@@ -195,7 +195,7 @@ def main() -> None:
     # into a JSON file that is compatible with OpenAPI schema specification
     if args.dump_schema:
         try:
-            schema_dumper.dump_schema("schema.json")
+            config_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)

src/utils/config_dumper.py

@@ -0,0 +1,23 @@
+"""Function to dump the configuration schema into OpenAPI-compatible format."""
+
+from models.config import Configuration
+from utils.openapi_schema_dumper import dump_openapi_schema
+
+
+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.
+    """
+    models = [Configuration]
+    dump_openapi_schema(models, filename)

src/utils/models_dumper.py

@@ -1,11 +1,7 @@
 """Function to dump the schema of all data models into OpenAPI-compatible format."""
 
-import json
-
-from pydantic.json_schema import models_json_schema
-
 import models.compaction as models_compaction
-from utils.json_schema_updater import recursive_update
+from utils.openapi_schema_dumper import dump_openapi_schema
 
 
 def dump_models(filename: str) -> None:
@@ -23,29 +19,5 @@ def dump_models(filename: str) -> None:
     ------
         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 [models_compaction.ConversationSummary]
-            ],
-            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)
+    models = [models_compaction.ConversationSummary]
+    dump_openapi_schema(models, filename)

src/utils/schema_dumper.py src/utils/openapi_schema_dumper.pysrc/utils/schema_dumper.py renamed to src/utils/openapi_schema_dumper.py

@@ -1,39 +1,30 @@
-"""Function to dump the configuration schema into OpenAPI-compatible format."""
+"""Utility function to dump schema with list of models into OpenAPI-compatible JSON format."""
 
 import json
 
 from pydantic.json_schema import models_json_schema
 
-from models.config import Configuration
 from utils.json_schema_updater import recursive_update
 
 
-def dump_schema(filename: str) -> None:
-    """Dump the configuration schema into OpenAPI-compatible JSON file.
+def dump_openapi_schema(models: list, filename: str) -> None:
+    """Write an OpenAPI-compatible JSON schema for the given models to a file.
 
     Parameters:
     ----------
+        - models: list - Pydantic model classes to include in the schema
         - 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}",
+            [(model, "validation") for model in models],
+            ref_template="`#/components/schemas/`{model}",
         )
-
-        # fix the schema
         schemas = recursive_update(schemas)
-
-        # add all required metadata
         openapi_schema = {
             "openapi": "3.0.0",
             "info": {

tests/unit/utils/test_schema_dumper.py tests/unit/utils/test_config_dumper.pytests/unit/utils/test_schema_dumper.py renamed to tests/unit/utils/test_config_dumper.py

@@ -1,9 +1,9 @@
-"""Unit tests for utils/schema_dumper module."""
+"""Unit tests for utils/config_dumper module."""
 
 from json import load
 from pathlib import Path
 
-from utils.schema_dumper import dump_schema
+from utils.config_dumper import dump_schema
 
 
 def test_dump_schema(tmpdir: Path) -> None: