Fixed types in models

road-core · Feb 9, 2025 · 774f696 · 774f696
1 parent ee2835d
commit 774f696
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/ols/app/models/models.py b/ols/app/models/models.py
@@ -760,8 +760,8 @@ def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(object_hook=self._decode_message, *args, **kwargs)
 
     def _decode_message(
-        self, dct: Dict[str, Any]
-    ) -> Union[HumanMessage, AIMessage, Dict[str, Any]]:
+        self, dct: dict[str, Any]
+    ) -> Union[HumanMessage, AIMessage, dict[str, Any]]:
         """Decode JSON dictionary into Message objects if applicable.
 
         Args:
@@ -771,6 +771,8 @@ def _decode_message(
             Union[HumanMessage, AIMessage, dict]: A Message object if the input
             dictionary represents a message, otherwise returns the original dictionary.
         """
+        message: BaseMessage
+
         if "type" in dct:
             if dct["type"] == "human":
                 message = HumanMessage(content=dct["content"])