This commit is contained in:
Chester Curme 2025-01-07 14:26:17 -05:00
parent c318ab6152
commit d18e1c67ff

View File

@ -21,6 +21,7 @@ from langchain_core.utils.function_calling import tool_example_to_messages
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic.v1 import BaseModel as BaseModelV1 from pydantic.v1 import BaseModel as BaseModelV1
from pydantic.v1 import Field as FieldV1 from pydantic.v1 import Field as FieldV1
from typing_extensions import Annotated, TypedDict
from langchain_tests.unit_tests.chat_models import ( from langchain_tests.unit_tests.chat_models import (
ChatModelTests, ChatModelTests,
@ -1293,6 +1294,7 @@ class ChatModelIntegrationTests(ChatModelTests):
if not self.has_tool_calling: if not self.has_tool_calling:
pytest.skip("Test requires tool calling.") pytest.skip("Test requires tool calling.")
# Pydantic
class Joke(BaseModel): class Joke(BaseModel):
"""Joke to tell user.""" """Joke to tell user."""
@ -1310,6 +1312,22 @@ class ChatModelIntegrationTests(ChatModelTests):
joke_result = chat.invoke("Give me a joke about cats, include the punchline.") joke_result = chat.invoke("Give me a joke about cats, include the punchline.")
assert isinstance(joke_result, Joke) 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: def test_json_mode(self, model: BaseChatModel) -> None:
"""Test structured output via `JSON mode. <https://python.langchain.com/docs/concepts/structured_outputs/#json-mode>`_ """Test structured output via `JSON mode. <https://python.langchain.com/docs/concepts/structured_outputs/#json-mode>`_