Implement Pydantic models for SIO events (#9)

* WIP implementation of neon-data-models SIO communication models

* Run pre-commit formatting

* WIP troubleshooting

* Confirmed functional CCAI performance with some bugs, likely related to client handling

* Clean up `sid` (shout_id)/`message_id` handling

* Remove extra log

* Remove unused import

* Update neon-data-models dependency spec

* Clean up temporary exception handling

* Clean up changes in `prompt.py`

* Clean up changes in `stt`

* Clean up changes in `tts`

* Clean up changes in `user_message`

* Update neon-data-models dependency spec

Author: NeonDaniel

Dockerfile

@@ -13,12 +13,10 @@ COPY . /app/
 WORKDIR /app
 
 RUN apt-get update \
-    && apt-get install -y \
-    && apt install build-essential -y \
-    && pip install --upgrade pip  \
-    && pip install wheel
+    && apt-get install build-essential git -y \
+    && pip install --no-cache-dir --upgrade pip wheel
 
-RUN pip install /app
+RUN pip install --no-cache-dir /app
 
 
 CMD ["pyklatchat-server"]

chat_server/blueprints/users.py

@@ -67,13 +67,16 @@ async def get_user(
     session_token = ""
     if user_id:
         user = MongoDocumentsAPI.USERS.get_user(user_id=user_id)
-        user.pop("password", None)
-        user.pop("date_created", None)
-        user.pop("tokens", None)
-        if session_data.user.user_id != user_id:
-            user.pop("roles", None)
-            user.pop("preferences", None)
-        LOG.info(f"Fetched user data (id={user_id}): {user}")
+        try:
+            user.pop("password", None)
+            user.pop("date_created", None)
+            user.pop("tokens", None)
+            if session_data.user.user_id != user_id:
+                user.pop("roles", None)
+                user.pop("preferences", None)
+            LOG.info(f"Fetched user data (id={user_id}): {user}")
+        except Exception as e:
+            raise RuntimeError(f"Failed to get user for user_id={user_id}") from e
     else:
         user = session_data.user
         session_token = session_data.session

chat_server/sio/handlers/prompt.py

@@ -26,11 +26,17 @@
 # NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 # SOFTWARE,  EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from time import time
-
 from klatchat_utils.database_utils.mongo_utils.queries import mongo_queries
-from klatchat_utils.database_utils.mongo_utils.queries.wrapper import MongoDocumentsAPI
+from klatchat_utils.database_utils.mongo_utils.queries.wrapper import (
+    MongoDocumentsAPI,
+)
 from neon_utils.logger import LOG
+from neon_data_models.models.api.klat.socketio import (
+    NewCcaiPrompt,
+    CcaiPromptCompleted,
+    GetPromptData,
+    PromptData,
+)
 from chat_server.sio.server import sio
 
 
@@ -40,33 +46,17 @@ async def new_prompt(sid, data):
     SIO event fired on new prompt data saving request
     :param sid: client session id
     :param data: user message data
-    Example:
-    ```
-        data = {'cid':'conversation id',
-                'promptID': 'id of related prompt',
-                'context': 'message context (optional)',
-                'timeCreated': 'timestamp on which message was created'
-                }
-    ```
     """
-    prompt_id = data["prompt_id"]
-    cid = data["cid"]
-    prompt_text = data["prompt_text"]
-    created_on = int(data.get("created_on") or time())
-    context = data.get("context") or {}
+    prompt = NewCcaiPrompt(**data)
+    LOG.debug(f"Creating new prompt: {prompt.prompt_text}")
     try:
-        formatted_data = {
-            "_id": prompt_id,
-            "cid": cid,
-            "is_completed": "0",
-            "data": {"prompt_text": prompt_text},
-            "context": context,
-            "created_on": created_on,
-        }
+        formatted_data = prompt.to_db_query()
         MongoDocumentsAPI.PROMPTS.add_item(data=formatted_data)
         await sio.emit("new_prompt_created", data=formatted_data)
     except Exception as ex:
-        LOG.error(f'Prompt "{prompt_id}" was not created due to exception - {ex}')
+        LOG.error(
+            f'Prompt "{prompt.prompt_id}" was not created due to exception - {ex}'
+        )
 
 
 @sio.event
@@ -76,17 +66,13 @@ async def prompt_completed(sid, data):
     :param sid: client session id
     :param data: user message data
     """
-    prompt_id = data["context"]["prompt"]["prompt_id"]
-
-    LOG.info(f"setting {prompt_id = } as completed")
-    MongoDocumentsAPI.PROMPTS.set_completed(
-        prompt_id=prompt_id, prompt_context=data["context"]
+    prompt = CcaiPromptCompleted(**data)
+    LOG.info(
+        f"setting prompt_id={prompt.prompt_id} as completed with: " f"{prompt.winner}"
     )
-    formatted_data = {
-        "winner": data["context"].get("winner", ""),
-        "prompt_id": prompt_id,
-    }
-    await sio.emit("set_prompt_completed", data=formatted_data)
+    MongoDocumentsAPI.PROMPTS.set_completed(**prompt.to_db_query())
+
+    await sio.emit("set_prompt_completed", data=prompt.model_dump())
 
 
 @sio.event
@@ -95,41 +81,29 @@ async def get_prompt_data(sid, data):
     SIO event fired getting prompt data request
     :param sid: client session id
     :param data: user message data
-    Example:
-    ```
-        data = {'userID': 'emitted user id',
-                'cid':'conversation id',
-                'promptID': 'id of related prompt'}
-    ```
     """
-    prompt_id = data.get("prompt_id")
-    _prompt_data = mongo_queries.fetch_prompt_data(
-        cid=data["cid"],
-        limit=data.get("limit", 5),
-        prompt_ids=[prompt_id],
-        fetch_user_data=True,
-    )
-    if prompt_id:
-        prompt_data = {
-            "_id": _prompt_data[0]["_id"],
-            "is_completed": _prompt_data[0].get("is_completed", "1"),
-            **_prompt_data[0].get("data", {}),
-        }
-    else:
-        prompt_data = []
-        for item in _prompt_data:
-            prompt_data.append(
-                {
-                    "_id": item["_id"],
-                    "created_on": item["created_on"],
-                    "is_completed": item.get("is_completed", "1"),
-                    **item["data"],
-                }
-            )
-    result = dict(
-        data=prompt_data,
-        receiver=data["nick"],
-        cid=data["cid"],
-        request_id=data["request_id"],
-    )
-    await sio.emit("prompt_data", data=result)
+    try:
+        requested_prompt_data = GetPromptData(**data)
+        _prompt_data = PromptData(
+            **mongo_queries.fetch_prompt_data(**requested_prompt_data.to_db_query())
+        )
+        if requested_prompt_data.prompt_id:
+            if isinstance(_prompt_data.data, list):
+                prompt_data = _prompt_data.data[0].model_dump()
+            else:
+                prompt_data = _prompt_data.data.model_dump()
+        else:
+            prompt_data = []
+            if isinstance(_prompt_data.data, list):
+                for item in _prompt_data:
+                    prompt_data.append(item.model_dump())
+        result = dict(
+            data=prompt_data,
+            receiver=requested_prompt_data.nick,
+            cid=requested_prompt_data.cid,
+            request_id=requested_prompt_data.request_id,
+        )
+        LOG.info(f"Emitting prompt_data: {result}")
+        await sio.emit("prompt_data", data=result)
+    except Exception as ex:
+        LOG.error(f"Failed to get prompt data due to exception - {ex}")

chat_server/sio/handlers/stt.py

@@ -28,35 +28,34 @@
 
 from klatchat_utils.database_utils.mongo_utils.queries.wrapper import MongoDocumentsAPI
 from neon_utils.logger import LOG
+from pydantic import ValidationError
 from chat_server.sio.server import sio
 from chat_server.sio.utils import emit_error
-from chat_server.utils.languages import LanguageSettings
+
+from neon_data_models.models.api.klat.socketio import GetSttResponse, GetSttRequest
 
 
 @sio.event
 async def stt_response(sid, data):
     """Handle STT Response from Observer"""
-    mq_context = data.get("context", {})
-    message_id = mq_context.get("message_id")
-    matching_shout = MongoDocumentsAPI.SHOUTS.get_item(item_id=message_id)
+    response = GetSttResponse(**data)
+    matching_shout = MongoDocumentsAPI.SHOUTS.get_item(item_id=response.sid)
     if not matching_shout:
         LOG.warning(
-            f"Skipping STT Response for message_id={message_id} - matching shout does not exist"
+            f"Skipping STT Response for sid={response.sid} - matching shout does not exist"
         )
     else:
         try:
-            message_text = data.get("transcript")
-            lang = LanguageSettings.to_system_lang(data["lang"])
             MongoDocumentsAPI.SHOUTS.save_stt_response(
-                shout_id=message_id, message_text=message_text, lang=lang
+                shout_id=response.sid,
+                message_text=response.transcript,
+                lang=response.lang,
             )
-            sid = mq_context.get("sid")
-            cid = mq_context.get("cid")
             response_data = {
-                "cid": cid,
-                "message_id": message_id,
-                "lang": lang,
-                "message_text": message_text,
+                "cid": response.cid,
+                "message_id": response.sid,
+                "lang": response.lang,
+                "message_text": response.transcript,
             }
             await sio.emit("incoming_stt", data=response_data, to=sid)
         except Exception as ex:
@@ -70,32 +69,18 @@ async def request_stt(sid, data):
 
     :param sid: client session id
     :param data: received tts request data
-    Example of tts request data:
-    ```
-        data = {
-                    'cid': (target conversation id)
-                    'message_id': (target message id),
-                    'audio_data':(target audio data base64 encoded),
-                    (optional) 'lang': (target message lang)
-               }
-    ```
     """
-    required_keys = ("message_id",)
-    if not all(key in list(data) for key in required_keys):
-        LOG.error(f"Missing one of the required keys - {required_keys}")
-    else:
-        cid = data.get("cid", "")
-        message_id = data.get("message_id", "")
-        # TODO: process received language
-        lang = "en"
-        # lang = data.get('lang', 'en')
-        if shout_data := MongoDocumentsAPI.SHOUTS.get_item(item_id=message_id):
-            message_transcript = shout_data.get("transcripts", {}).get(lang)
+    try:
+        request = GetSttRequest(sid=sid, **data)
+        # TODO: Identify reason for this language patch
+        request.lang = "en"
+        if shout_data := MongoDocumentsAPI.SHOUTS.get_item(item_id=request.sid):
+            message_transcript = shout_data.get("transcripts", {}).get(request.lang)
             if message_transcript:
                 response_data = {
-                    "cid": cid,
-                    "message_id": message_id,
-                    "lang": lang,
+                    "cid": request.cid,
+                    "message_id": request.sid,
+                    "lang": request.lang,
                     "message_text": message_transcript,
                 }
                 return await sio.emit("incoming_stt", data=response_data, to=sid)
@@ -105,16 +90,10 @@ async def request_stt(sid, data):
                 return await emit_error(message=err_msg, sids=[sid])
         audio_data = data.get(
             "audio_data"
-        ) or MongoDocumentsAPI.SHOUTS.fetch_audio_data(message_id=message_id)
+        ) or MongoDocumentsAPI.SHOUTS.fetch_audio_data(sid=request.sid)
         if not audio_data:
             LOG.error("Failed to fetch audio data")
         else:
-            lang = LanguageSettings.to_neon_lang(lang)
-            formatted_data = {
-                "cid": cid,
-                "sid": sid,
-                "message_id": message_id,
-                "audio_data": audio_data,
-                "lang": lang,
-            }
-            await sio.emit("get_stt", data=formatted_data)
+            await sio.emit("get_stt", data=request.model_dump())
+    except ValidationError:
+        LOG.exception(f"Invalid STT request data - {data}")

chat_server/sio/handlers/translation.py

@@ -35,6 +35,10 @@
 from chat_server.sio.server import sio
 from chat_server.utils.cache_utils import CacheFactory
 
+# TODO: The backend service supporting these endpoints is no longer functional.
+# These need to be updated to implement appropriate data models once the
+# backend service is re-implemented.
+
 
 @sio.event
 async def request_translate(sid, data):