Skip to content

Commit

Permalink
fix type check
Browse files Browse the repository at this point in the history
Signed-off-by: Stephanie <[email protected]>
  • Loading branch information
yangcao77 committed Jan 22, 2025
1 parent 837b7fe commit 97fabd9
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions ols/app/models/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import json
from collections import OrderedDict
from typing import Optional, Self
from typing import Any, Dict, Optional, Self, Union

from langchain.llms.base import LLM
from pydantic import BaseModel, field_validator, model_validator
Expand Down Expand Up @@ -713,7 +713,7 @@ def cache_entries_to_history(
class MessageEncoder(json.JSONEncoder):
"""Convert Message objects to serializable dictionaries."""

def default(self, o):
def default(self, o: Any) -> Union[dict, Any]:
"""Convert a Message object into a serializable dictionary.
This method is called when an object cannot be serialized by default
Expand Down Expand Up @@ -756,11 +756,11 @@ class MessageDecoder(json.JSONDecoder):
HumanMessage(content="Hello", ...)
"""

def __init__(self, *args, **kwargs):
def __init__(self, *args: Any, **kwargs: Any):
"""Initialize the MessageDecoder with custom object hook."""
super().__init__(object_hook=self._decode_message, *args, **kwargs)

def _decode_message(self, dct):
def _decode_message(self, dct: Dict[str, Any]) -> Union[HumanMessage, AIMessage, Dict[str, Any]]:
"""Decode JSON dictionary into Message objects if applicable.
Args:
Expand Down

0 comments on commit 97fabd9

Please sign in to comment.