diff --git a/libs/partners/mongodb/tests/utils.py b/libs/partners/mongodb/tests/utils.py index 674bd2c43fd0e..d62b2025f48ea 100644 --- a/libs/partners/mongodb/tests/utils.py +++ b/libs/partners/mongodb/tests/utils.py @@ -17,7 +17,7 @@ BaseMessage, ) from langchain_core.outputs import ChatGeneration, ChatResult -from pydantic import validator +from pydantic import model_validator from pymongo.collection import Collection from pymongo.results import DeleteResult, InsertManyResult @@ -134,15 +134,14 @@ class FakeLLM(LLM): sequential_responses: Optional[bool] = False response_index: int = 0 - @validator("queries", always=True) - def check_queries_required( - cls, queries: Optional[Mapping], values: Mapping[str, Any] - ) -> Optional[Mapping]: - if values.get("sequential_response") and not queries: + @model_validator(mode="before") + @classmethod + def check_queries_required(cls, values: dict) -> dict: + if values.get("sequential_response") and not values.get("queries"): raise ValueError( "queries is required when sequential_response is set to True" ) - return queries + return values def get_num_tokens(self, text: str) -> int: """Return number of tokens."""