Merge pull request #2096 from tisnik/lcore-2755-model-dumper

LCORE-2755: Model dumper

Author: tisnik

README.md

@@ -858,6 +858,7 @@ options:
   -d, --dump-configuration
                         dump actual configuration into JSON file and quit
   -s, --dump-schema     dump configuration schema into OpenAPI-compatible file and quit
+  -m, --dump-models     dump schemas for all models into OpenAPI-compatible file and quit
   -c, --config CONFIG_FILE
                         path to configuration file (default: lightspeed-stack.yaml)
   --synthesized-config-output SYNTHESIZED_CONFIG_OUTPUT

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 schema_dumper
+from utils import models_dumper, schema_dumper
 
 setup_logging()
 logger = get_logger(__name__)
@@ -61,6 +61,14 @@ def create_argument_parser() -> ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "-m",
+        "--dump-models",
+        dest="dump_models",
+        help="dump schemas for all models into OpenAPI-compatible file and quit",
+        action="store_true",
+        default=False,
+    )
     parser.add_argument(
         "-c",
         "--config",
@@ -105,6 +113,7 @@ def create_argument_parser() -> ArgumentParser:
     return parser
 
 
+# pylint: disable=too-many-branches, too-many-statements
 def main() -> None:
     """Entry point to the web service.
 
@@ -193,6 +202,17 @@ def main() -> None:
             raise SystemExit(1) from e
         return
 
+    # -m or --dump-models CLI flags are used to dump schema for all models
+    # into a JSON file that is compatible with OpenAPI schema specification
+    if args.dump_models:
+        try:
+            models_dumper.dump_models("models.json")
+            logger.info("Schema for all models dumped to models.json")
+        except Exception as e:
+            logger.error("Failed to dump schema for models: %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[constants.CONFIG_PATH_ENV_VAR] = args.config_file

src/utils/models_dumper.py

@@ -0,0 +1,51 @@
+"""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
+
+
+def dump_models(filename: str) -> None:
+    """Dump the schema of all models 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 [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)

tests/unit/utils/test_models_dumper.py

@@ -0,0 +1,92 @@
+"""Unit tests for utils/models_dumper module."""
+
+from json import load
+from pathlib import Path
+
+from utils.models_dumper import dump_models
+
+
+def test_dump_models(tmpdir: Path) -> None:
+    """Test that models can be dump into a JSON file.
+
+    An example of schema dump:
+    {
+    "openapi": "3.0.0",
+    "info": {
+        "title": "Lightspeed Core Stack",
+        "version": "0.3.0"
+    },
+    "components": {
+        "schemas": {
+            "ConversationSummary": {
+                "description": "A single compaction-produced summary chunk.\n\nAttributes:\n",
+                "properties": {
+                    "summary_text": {
+                        "description": "Natural-language summary produced by the ...",
+                        "title": "Summary text",
+                        "type": "string"
+                    },
+                    "summarized_through_turn": {
+                        "description": "Running total of conversation items consumed by this ...",
+                        "minimum": 0,
+                        "title": "Summarized through turn",
+                        "type": "integer"
+                    },
+                    "token_count": {
+                        "description": "Number of tokens in summary_text.",
+                        "minimum": 0,
+                        "title": "Token count",
+                        "type": "integer"
+                    },
+                    "created_at": {
+                        "description": "ISO 8601 timestamp recording when this summary ...",
+                        "title": "Created at",
+                        "type": "string"
+                    },
+                    "model_used": {
+                        "description": "Fully-qualified model identifier used for the ...",
+                        "title": "Model used",
+                        "type": "string"
+                    }
+                },
+                "required": [
+                    "summary_text",
+                    "summarized_through_turn",
+                    "token_count",
+                    "created_at",
+                    "model_used"
+                ],
+                "title": "ConversationSummary",
+                "type": "object"
+            }
+        }
+    },
+    "paths": {}
+    }
+    """
+    filename = tmpdir / "foo.json"
+    dump_models(str(filename))
+
+    with open(filename, "r", encoding="utf-8") as fin:
+        # schema should be stored in JSON format
+        content = load(fin)
+        assert content is not None
+
+        # top-level keys test
+        keys = ("openapi", "info", "components", "paths")
+        for key in keys:
+            assert key in content
+
+        # components should be top-level node
+        components = content["components"]
+        assert components is not None
+
+        # schemas should be a node stored inside components node
+        assert "schemas" in components
+        schemas = components["schemas"]
+        assert schemas is not None
+
+        # list of schemas expected in a dump
+        expected_schemas = ("ConversationSummary",)
+        for expected_schema in expected_schemas:
+            assert expected_schema in schemas