From 99adcdb1c916cc30e84f25a194afa1e427db4b5c Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 11 Oct 2023 15:06:42 -0400 Subject: [PATCH] Add dedicated `type` attribute to be used solely for serialization purposes (#11585) Adds standard `type` field for all messages that will be serialized/validated by pydantic. * The presence of `type` makes it easier for developers consuming schemas to write client code to serialize/deserialize. * In LangServe `type` will be used for both validation and will appear in the generated openapi specs --- libs/langchain/langchain/prompts/base.py | 3 +- libs/langchain/langchain/prompts/chat.py | 3 + libs/langchain/langchain/schema/document.py | 3 +- libs/langchain/langchain/schema/messages.py | 17 +- .../runnable/__snapshots__/test_runnable.ambr | 380 +++++------------- .../schema/runnable/test_runnable.py | 54 ++- .../langchain/tests/unit_tests/test_schema.py | 35 +- 7 files changed, 152 insertions(+), 343 deletions(-) diff --git a/libs/langchain/langchain/prompts/base.py b/libs/langchain/langchain/prompts/base.py index df2e558cd5..f2cbd6dad4 100644 --- a/libs/langchain/langchain/prompts/base.py +++ b/libs/langchain/langchain/prompts/base.py @@ -3,7 +3,7 @@ from __future__ import annotations import warnings from abc import ABC -from typing import Any, Callable, Dict, List, Set +from typing import Any, Callable, Dict, List, Literal, Set from langchain.schema.messages import BaseMessage, HumanMessage from langchain.schema.prompt import PromptValue @@ -104,6 +104,7 @@ class StringPromptValue(PromptValue): text: str """Prompt text.""" + type: Literal["StringPromptValue"] = "StringPromptValue" def to_string(self) -> str: """Return prompt as string.""" diff --git a/libs/langchain/langchain/prompts/chat.py b/libs/langchain/langchain/prompts/chat.py index b93fb028bd..8c301c109a 100644 --- a/libs/langchain/langchain/prompts/chat.py +++ b/libs/langchain/langchain/prompts/chat.py @@ -8,6 +8,7 @@ from typing import ( Callable, Dict, List, + Literal, Sequence, Set, Tuple, @@ -299,6 +300,8 @@ class ChatPromptValueConcrete(ChatPromptValue): messages: Sequence[AnyMessage] + type: Literal["ChatPromptValueConcrete"] = "ChatPromptValueConcrete" + class BaseChatPromptTemplate(BasePromptTemplate, ABC): """Base class for chat prompt templates.""" diff --git a/libs/langchain/langchain/schema/document.py b/libs/langchain/langchain/schema/document.py index acb46dd1d2..c552ebda7d 100644 --- a/libs/langchain/langchain/schema/document.py +++ b/libs/langchain/langchain/schema/document.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio from abc import ABC, abstractmethod from functools import partial -from typing import Any, Sequence +from typing import Any, Literal, Sequence from langchain.load.serializable import Serializable from langchain.pydantic_v1 import Field @@ -18,6 +18,7 @@ class Document(Serializable): """Arbitrary metadata about the page content (e.g., source, relationships to other documents, etc.). """ + type: Literal["Document"] = "Document" @classmethod def is_lc_serializable(cls) -> bool: diff --git a/libs/langchain/langchain/schema/messages.py b/libs/langchain/langchain/schema/messages.py index f2e33b0248..1edb5d1bd8 100644 --- a/libs/langchain/langchain/schema/messages.py +++ b/libs/langchain/langchain/schema/messages.py @@ -149,7 +149,6 @@ class HumanMessage(BaseMessage): """ type: Literal["human"] = "human" - is_chunk: Literal[False] = False HumanMessage.update_forward_refs() @@ -161,7 +160,7 @@ class HumanMessageChunk(HumanMessage, BaseMessageChunk): # Ignoring mypy re-assignment here since we're overriding the value # to make sure that the chunk variant can be discriminated from the # non-chunk variant. - is_chunk: Literal[True] = True # type: ignore[assignment] + type: Literal["HumanMessageChunk"] = "HumanMessageChunk" # type: ignore[assignment] # noqa: E501 class AIMessage(BaseMessage): @@ -173,7 +172,6 @@ class AIMessage(BaseMessage): """ type: Literal["ai"] = "ai" - is_chunk: Literal[False] = False AIMessage.update_forward_refs() @@ -185,7 +183,7 @@ class AIMessageChunk(AIMessage, BaseMessageChunk): # Ignoring mypy re-assignment here since we're overriding the value # to make sure that the chunk variant can be discriminated from the # non-chunk variant. - is_chunk: Literal[True] = True # type: ignore[assignment] + type: Literal["AIMessageChunk"] = "AIMessageChunk" # type: ignore[assignment] # noqa: E501 def __add__(self, other: Any) -> BaseMessageChunk: # type: ignore if isinstance(other, AIMessageChunk): @@ -211,7 +209,6 @@ class SystemMessage(BaseMessage): """ type: Literal["system"] = "system" - is_chunk: Literal[False] = False SystemMessage.update_forward_refs() @@ -223,7 +220,7 @@ class SystemMessageChunk(SystemMessage, BaseMessageChunk): # Ignoring mypy re-assignment here since we're overriding the value # to make sure that the chunk variant can be discriminated from the # non-chunk variant. - is_chunk: Literal[True] = True # type: ignore[assignment] + type: Literal["SystemMessageChunk"] = "SystemMessageChunk" # type: ignore[assignment] # noqa: E501 class FunctionMessage(BaseMessage): @@ -233,7 +230,6 @@ class FunctionMessage(BaseMessage): """The name of the function that was executed.""" type: Literal["function"] = "function" - is_chunk: Literal[False] = False FunctionMessage.update_forward_refs() @@ -245,7 +241,9 @@ class FunctionMessageChunk(FunctionMessage, BaseMessageChunk): # Ignoring mypy re-assignment here since we're overriding the value # to make sure that the chunk variant can be discriminated from the # non-chunk variant. - is_chunk: Literal[True] = True # type: ignore[assignment] + type: Literal[ + "FunctionMessageChunk" + ] = "FunctionMessageChunk" # type: ignore[assignment] def __add__(self, other: Any) -> BaseMessageChunk: # type: ignore if isinstance(other, FunctionMessageChunk): @@ -272,7 +270,6 @@ class ChatMessage(BaseMessage): """The speaker / role of the Message.""" type: Literal["chat"] = "chat" - is_chunk: Literal[False] = False ChatMessage.update_forward_refs() @@ -284,7 +281,7 @@ class ChatMessageChunk(ChatMessage, BaseMessageChunk): # Ignoring mypy re-assignment here since we're overriding the value # to make sure that the chunk variant can be discriminated from the # non-chunk variant. - is_chunk: Literal[True] = True # type: ignore[assignment] + type: Literal["ChatMessageChunk"] = "ChatMessageChunk" # type: ignore def __add__(self, other: Any) -> BaseMessageChunk: # type: ignore if isinstance(other, ChatMessageChunk): diff --git a/libs/langchain/tests/unit_tests/schema/runnable/__snapshots__/test_runnable.ambr b/libs/langchain/tests/unit_tests/schema/runnable/__snapshots__/test_runnable.ambr index 15d2e82970..fb9d63ee91 100644 --- a/libs/langchain/tests/unit_tests/schema/runnable/__snapshots__/test_runnable.ambr +++ b/libs/langchain/tests/unit_tests/schema/runnable/__snapshots__/test_runnable.ambr @@ -1693,14 +1693,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'ai', 'enum': list([ @@ -1727,14 +1719,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'role': dict({ 'title': 'Role', 'type': 'string', @@ -1784,6 +1768,14 @@ 'title': 'Messages', 'type': 'array', }), + 'type': dict({ + 'default': 'ChatPromptValueConcrete', + 'enum': list([ + 'ChatPromptValueConcrete', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'messages', @@ -1802,14 +1794,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'name': dict({ 'title': 'Name', 'type': 'string', @@ -1846,14 +1830,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'human', 'enum': list([ @@ -1876,6 +1852,14 @@ 'title': 'Text', 'type': 'string', }), + 'type': dict({ + 'default': 'StringPromptValue', + 'enum': list([ + 'StringPromptValue', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'text', @@ -1897,14 +1881,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'system', 'enum': list([ @@ -1976,14 +1952,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'ai', 'enum': list([ @@ -2010,14 +1978,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'role': dict({ 'title': 'Role', 'type': 'string', @@ -2067,6 +2027,14 @@ 'title': 'Messages', 'type': 'array', }), + 'type': dict({ + 'default': 'ChatPromptValueConcrete', + 'enum': list([ + 'ChatPromptValueConcrete', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'messages', @@ -2085,14 +2053,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'name': dict({ 'title': 'Name', 'type': 'string', @@ -2129,14 +2089,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'human', 'enum': list([ @@ -2159,6 +2111,14 @@ 'title': 'Text', 'type': 'string', }), + 'type': dict({ + 'default': 'StringPromptValue', + 'enum': list([ + 'StringPromptValue', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'text', @@ -2180,14 +2140,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'system', 'enum': list([ @@ -2243,18 +2195,10 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': True, - 'enum': list([ - True, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ - 'default': 'ai', + 'default': 'AIMessageChunk', 'enum': list([ - 'ai', + 'AIMessageChunk', ]), 'title': 'Type', 'type': 'string', @@ -2277,22 +2221,14 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': True, - 'enum': list([ - True, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'role': dict({ 'title': 'Role', 'type': 'string', }), 'type': dict({ - 'default': 'chat', + 'default': 'ChatMessageChunk', 'enum': list([ - 'chat', + 'ChatMessageChunk', ]), 'title': 'Type', 'type': 'string', @@ -2316,22 +2252,14 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': True, - 'enum': list([ - True, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'name': dict({ 'title': 'Name', 'type': 'string', }), 'type': dict({ - 'default': 'function', + 'default': 'FunctionMessageChunk', 'enum': list([ - 'function', + 'FunctionMessageChunk', ]), 'title': 'Type', 'type': 'string', @@ -2360,18 +2288,10 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': True, - 'enum': list([ - True, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ - 'default': 'human', + 'default': 'HumanMessageChunk', 'enum': list([ - 'human', + 'HumanMessageChunk', ]), 'title': 'Type', 'type': 'string', @@ -2394,18 +2314,10 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': True, - 'enum': list([ - True, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ - 'default': 'system', + 'default': 'SystemMessageChunk', 'enum': list([ - 'system', + 'SystemMessageChunk', ]), 'title': 'Type', 'type': 'string', @@ -2448,14 +2360,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'ai', 'enum': list([ @@ -2482,14 +2386,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'role': dict({ 'title': 'Role', 'type': 'string', @@ -2539,6 +2435,14 @@ 'title': 'Messages', 'type': 'array', }), + 'type': dict({ + 'default': 'ChatPromptValueConcrete', + 'enum': list([ + 'ChatPromptValueConcrete', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'messages', @@ -2557,14 +2461,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'name': dict({ 'title': 'Name', 'type': 'string', @@ -2601,14 +2497,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'human', 'enum': list([ @@ -2631,6 +2519,14 @@ 'title': 'Text', 'type': 'string', }), + 'type': dict({ + 'default': 'StringPromptValue', + 'enum': list([ + 'StringPromptValue', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'text', @@ -2652,14 +2548,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'system', 'enum': list([ @@ -2706,14 +2594,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'ai', 'enum': list([ @@ -2740,14 +2620,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'role': dict({ 'title': 'Role', 'type': 'string', @@ -2797,6 +2669,14 @@ 'title': 'Messages', 'type': 'array', }), + 'type': dict({ + 'default': 'ChatPromptValueConcrete', + 'enum': list([ + 'ChatPromptValueConcrete', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'messages', @@ -2815,14 +2695,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'name': dict({ 'title': 'Name', 'type': 'string', @@ -2859,14 +2731,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'human', 'enum': list([ @@ -2889,6 +2753,14 @@ 'title': 'Text', 'type': 'string', }), + 'type': dict({ + 'default': 'StringPromptValue', + 'enum': list([ + 'StringPromptValue', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'text', @@ -2910,14 +2782,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'system', 'enum': list([ @@ -2956,14 +2820,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'ai', 'enum': list([ @@ -2990,14 +2846,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'role': dict({ 'title': 'Role', 'type': 'string', @@ -3047,6 +2895,14 @@ 'title': 'Messages', 'type': 'array', }), + 'type': dict({ + 'default': 'ChatPromptValueConcrete', + 'enum': list([ + 'ChatPromptValueConcrete', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'messages', @@ -3065,14 +2921,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'name': dict({ 'title': 'Name', 'type': 'string', @@ -3109,14 +2957,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'human', 'enum': list([ @@ -3150,6 +2990,14 @@ 'title': 'Text', 'type': 'string', }), + 'type': dict({ + 'default': 'StringPromptValue', + 'enum': list([ + 'StringPromptValue', + ]), + 'title': 'Type', + 'type': 'string', + }), }), 'required': list([ 'text', @@ -3171,14 +3019,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'system', 'enum': list([ @@ -3241,14 +3081,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'ai', 'enum': list([ @@ -3275,14 +3107,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'role': dict({ 'title': 'Role', 'type': 'string', @@ -3314,14 +3138,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'name': dict({ 'title': 'Name', 'type': 'string', @@ -3358,14 +3174,6 @@ 'title': 'Example', 'type': 'boolean', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'human', 'enum': list([ @@ -3395,14 +3203,6 @@ 'title': 'Content', 'type': 'string', }), - 'is_chunk': dict({ - 'default': False, - 'enum': list([ - False, - ]), - 'title': 'Is Chunk', - 'type': 'boolean', - }), 'type': dict({ 'default': 'system', 'enum': list([ diff --git a/libs/langchain/tests/unit_tests/schema/runnable/test_runnable.py b/libs/langchain/tests/unit_tests/schema/runnable/test_runnable.py index 45cc12ba95..20e316dbb0 100644 --- a/libs/langchain/tests/unit_tests/schema/runnable/test_runnable.py +++ b/libs/langchain/tests/unit_tests/schema/runnable/test_runnable.py @@ -227,6 +227,12 @@ def test_schemas(snapshot: SnapshotAssertion) -> None: "properties": { "page_content": {"title": "Page Content", "type": "string"}, "metadata": {"title": "Metadata", "type": "object"}, + "type": { + "title": "Type", + "enum": ["Document"], + "default": "Document", + "type": "string", + }, }, "required": ["page_content"], } @@ -293,12 +299,6 @@ def test_schemas(snapshot: SnapshotAssertion) -> None: "default": False, "type": "boolean", }, - "is_chunk": { - "title": "Is Chunk", - "default": False, - "enum": [False], - "type": "boolean", - }, }, "required": ["content"], }, @@ -323,12 +323,6 @@ def test_schemas(snapshot: SnapshotAssertion) -> None: "default": False, "type": "boolean", }, - "is_chunk": { - "title": "Is Chunk", - "default": False, - "enum": [False], - "type": "boolean", - }, }, "required": ["content"], }, @@ -349,12 +343,6 @@ def test_schemas(snapshot: SnapshotAssertion) -> None: "type": "string", }, "role": {"title": "Role", "type": "string"}, - "is_chunk": { - "title": "Is Chunk", - "default": False, - "enum": [False], - "type": "boolean", - }, }, "required": ["content", "role"], }, @@ -374,12 +362,6 @@ def test_schemas(snapshot: SnapshotAssertion) -> None: "enum": ["system"], "type": "string", }, - "is_chunk": { - "title": "Is Chunk", - "default": False, - "enum": [False], - "type": "boolean", - }, }, "required": ["content"], }, @@ -400,12 +382,6 @@ def test_schemas(snapshot: SnapshotAssertion) -> None: "type": "string", }, "name": {"title": "Name", "type": "string"}, - "is_chunk": { - "title": "Is Chunk", - "default": False, - "enum": [False], - "type": "boolean", - }, }, "required": ["content", "name"], }, @@ -634,6 +610,12 @@ def test_schema_chains() -> None: "properties": { "page_content": {"title": "Page Content", "type": "string"}, "metadata": {"title": "Metadata", "type": "object"}, + "type": { + "title": "Type", + "type": "string", + "enum": ["Document"], + "default": "Document", + }, }, "required": ["page_content"], } @@ -667,6 +649,12 @@ def test_schema_chains() -> None: "properties": { "page_content": {"title": "Page Content", "type": "string"}, "metadata": {"title": "Metadata", "type": "object"}, + "type": { + "title": "Type", + "type": "string", + "enum": ["Document"], + "default": "Document", + }, }, "required": ["page_content"], } @@ -705,6 +693,12 @@ def test_schema_chains() -> None: "properties": { "page_content": {"title": "Page Content", "type": "string"}, "metadata": {"title": "Metadata", "type": "object"}, + "type": { + "title": "Type", + "type": "string", + "enum": ["Document"], + "default": "Document", + }, }, "required": ["page_content"], } diff --git a/libs/langchain/tests/unit_tests/test_schema.py b/libs/langchain/tests/unit_tests/test_schema.py index 4b72ddde1c..629f5bcbeb 100644 --- a/libs/langchain/tests/unit_tests/test_schema.py +++ b/libs/langchain/tests/unit_tests/test_schema.py @@ -3,7 +3,10 @@ import unittest from typing import Union +from langchain.prompts.base import StringPromptValue +from langchain.prompts.chat import ChatPromptValueConcrete from langchain.pydantic_v1 import BaseModel +from langchain.schema import Document from langchain.schema.messages import ( AIMessage, AIMessageChunk, @@ -81,24 +84,29 @@ def test_multiple_msg() -> None: assert messages_from_dict(messages_to_dict(msgs)) == msgs -def test_distinguish_messages() -> None: - """Test that pydantic is able to discriminate between similar looking messages.""" +def test_serialization_of_wellknown_objects() -> None: + """Test that pydantic is able to serialize and deserialize well known objects.""" + + class WellKnownLCObject(BaseModel): + """A well known LangChain object.""" - class WellKnownTypes(BaseModel): __root__: Union[ + Document, HumanMessage, - AIMessage, SystemMessage, + ChatMessage, FunctionMessage, + AIMessage, HumanMessageChunk, - AIMessageChunk, SystemMessageChunk, - FunctionMessageChunk, ChatMessageChunk, - ChatMessage, + FunctionMessageChunk, + AIMessageChunk, + StringPromptValue, + ChatPromptValueConcrete, ] - messages = [ + lc_objects = [ HumanMessage(content="human"), HumanMessageChunk(content="human"), AIMessage(content="ai"), @@ -121,8 +129,13 @@ def test_distinguish_messages() -> None: role="human", content="human", ), + StringPromptValue(text="hello"), + ChatPromptValueConcrete(messages=[HumanMessage(content="human")]), + Document(page_content="hello"), ] - for msg in messages: - obj1 = WellKnownTypes.parse_obj(msg.dict()) - assert type(obj1.__root__) == type(msg), f"failed for {type(msg)}" + for lc_object in lc_objects: + d = lc_object.dict() + assert "type" in d, f"Missing key `type` for {type(lc_object)}" + obj1 = WellKnownLCObject.parse_obj(d) + assert type(obj1.__root__) == type(lc_object), f"failed for {type(lc_object)}"