From d18e1c67ff4cf1f3cb7550cbcf0ce014cfee83d5 Mon Sep 17 00:00:00 2001 From: Chester Curme Date: Tue, 7 Jan 2025 14:26:17 -0500 Subject: [PATCH] add test --- .../integration_tests/chat_models.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/libs/standard-tests/langchain_tests/integration_tests/chat_models.py b/libs/standard-tests/langchain_tests/integration_tests/chat_models.py index f5691350044..26710f03cdc 100644 --- a/libs/standard-tests/langchain_tests/integration_tests/chat_models.py +++ b/libs/standard-tests/langchain_tests/integration_tests/chat_models.py @@ -21,6 +21,7 @@ from langchain_core.utils.function_calling import tool_example_to_messages from pydantic import BaseModel, Field from pydantic.v1 import BaseModel as BaseModelV1 from pydantic.v1 import Field as FieldV1 +from typing_extensions import Annotated, TypedDict from langchain_tests.unit_tests.chat_models import ( ChatModelTests, @@ -1293,6 +1294,7 @@ class ChatModelIntegrationTests(ChatModelTests): if not self.has_tool_calling: pytest.skip("Test requires tool calling.") + # Pydantic class Joke(BaseModel): """Joke to tell user.""" @@ -1310,6 +1312,22 @@ class ChatModelIntegrationTests(ChatModelTests): joke_result = chat.invoke("Give me a joke about cats, include the punchline.") assert isinstance(joke_result, Joke) + # Schema + chat = model.with_structured_output(Joke.model_json_schema()) + result = chat.invoke("Tell me a joke about cats.") + assert isinstance(result, dict) + + # TypedDict + class JokeDict(TypedDict): + """Joke to tell user.""" + + setup: Annotated[str, ..., "question to set up a joke"] + punchline: Annotated[Optional[str], None, "answer to resolve the joke"] + + chat = model.with_structured_output(JokeDict) + result = chat.invoke("Tell me a joke about cats.") + assert isinstance(result, dict) + def test_json_mode(self, model: BaseChatModel) -> None: """Test structured output via `JSON mode. `_