From 1dea16ab5862b0ab4b6cfda63c9d2dc85c4ac8fa Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Thu, 9 Jan 2025 17:09:46 -0800 Subject: [PATCH] Override `_PromptData.model_dump` to dump by alias by default Add unit test case to validate PromptData serialization --- neon_data_models/models/api/klat/socketio.py | 6 +++- tests/models/api/test_klat.py | 37 ++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) create mode 100644 tests/models/api/test_klat.py diff --git a/neon_data_models/models/api/klat/socketio.py b/neon_data_models/models/api/klat/socketio.py index e7720d0..1f17265 100644 --- a/neon_data_models/models/api/klat/socketio.py +++ b/neon_data_models/models/api/klat/socketio.py @@ -117,7 +117,7 @@ class GetPromptData(BaseModel): class PromptData(BaseModel): class _PromptData(BaseModel): - _id: str = Field(description="Unique ID for the prompt") + id: str = Field(alias="_id", description="Unique ID for the prompt") is_completed: Literal['0', '1'] = Field( description="'1' if a response to the prompt has been determined") proposed_responses: Dict[str, str] = Field( @@ -133,6 +133,10 @@ class _PromptData(BaseModel): default=[], description="List of subminds that participated in this prompt") + def model_dump(self, *args, **kwargs): + kwargs.setdefault('by_alias', True) + return BaseModel.model_dump(self, *args, **kwargs) + data: Union[_PromptData, List[_PromptData]] = Field(description="Prompt data") receiver: str = Field(description="Nickname of user requesting prompt data") cid: str = Field(description="Conversation ID associated with the prompt") diff --git a/tests/models/api/test_klat.py b/tests/models/api/test_klat.py new file mode 100644 index 0000000..f74bea0 --- /dev/null +++ b/tests/models/api/test_klat.py @@ -0,0 +1,37 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Development System +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2024 Neongecko.com Inc. +# BSD-3 +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from unittest import TestCase + + +class TestKlat(TestCase): + def test_prompt_data(self): + from neon_data_models.models.api.klat.socketio import PromptData + test_object = PromptData._PromptData(_id="test_id", + is_completed='1') + self.assertIsInstance(test_object, PromptData._PromptData) + self.assertEqual(test_object, + PromptData._PromptData(**test_object.model_dump()))