core[minor]: add name to basemessage (#17539)

Adds an optional name param to our base message to support passing names
into LLMs.

OpenAI supports having a name on anything except tool message now
(system, ai, user/human).
This commit is contained in:
Erick Friis 2024-02-14 12:21:59 -08:00 committed by GitHub
parent 916332ef5b
commit 86d3e42853
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 227 additions and 18 deletions

View File

@ -1,6 +1,6 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Dict, List, Sequence, Union
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
from langchain_core.load.serializable import Serializable
from langchain_core.pydantic_v1 import Extra, Field
@ -25,6 +25,8 @@ class BaseMessage(Serializable):
type: str
name: Optional[str] = None
class Config:
extra = Extra.allow
@ -53,6 +55,8 @@ class BaseMessage(Serializable):
def pretty_repr(self, html: bool = False) -> str:
title = get_msg_title_repr(self.type.title() + " Message", bold=html)
# TODO: handle non-string content.
if self.name is not None:
title += f"\nName: {self.name}"
return f"{title}\n\n{self.content}"
def pretty_print(self) -> None:

View File

@ -1718,6 +1718,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'ai',
'enum': list([
@ -1761,6 +1765,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'role': dict({
'title': 'Role',
'type': 'string',
@ -1909,6 +1917,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'human',
'enum': list([
@ -1977,6 +1989,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'system',
'enum': list([
@ -2020,6 +2036,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'tool_call_id': dict({
'title': 'Tool Call Id',
'type': 'string',
@ -2116,6 +2136,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'ai',
'enum': list([
@ -2159,6 +2183,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'role': dict({
'title': 'Role',
'type': 'string',
@ -2307,6 +2335,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'human',
'enum': list([
@ -2375,6 +2407,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'system',
'enum': list([
@ -2418,6 +2454,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'tool_call_id': dict({
'title': 'Tool Call Id',
'type': 'string',
@ -2498,6 +2538,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'ai',
'enum': list([
@ -2541,6 +2585,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'role': dict({
'title': 'Role',
'type': 'string',
@ -2642,6 +2690,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'human',
'enum': list([
@ -2688,6 +2740,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'system',
'enum': list([
@ -2731,6 +2787,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'tool_call_id': dict({
'title': 'Tool Call Id',
'type': 'string',
@ -2799,6 +2859,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'ai',
'enum': list([
@ -2842,6 +2906,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'role': dict({
'title': 'Role',
'type': 'string',
@ -2990,6 +3058,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'human',
'enum': list([
@ -3058,6 +3130,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'system',
'enum': list([
@ -3101,6 +3177,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'tool_call_id': dict({
'title': 'Tool Call Id',
'type': 'string',
@ -3169,6 +3249,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'ai',
'enum': list([
@ -3212,6 +3296,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'role': dict({
'title': 'Role',
'type': 'string',
@ -3360,6 +3448,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'human',
'enum': list([
@ -3428,6 +3520,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'system',
'enum': list([
@ -3471,6 +3567,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'tool_call_id': dict({
'title': 'Tool Call Id',
'type': 'string',
@ -3531,6 +3631,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'ai',
'enum': list([
@ -3574,6 +3678,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'role': dict({
'title': 'Role',
'type': 'string',
@ -3722,6 +3830,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'human',
'enum': list([
@ -3801,6 +3913,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'system',
'enum': list([
@ -3844,6 +3960,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'tool_call_id': dict({
'title': 'Tool Call Id',
'type': 'string',
@ -3931,6 +4051,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'ai',
'enum': list([
@ -3974,6 +4098,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'role': dict({
'title': 'Role',
'type': 'string',
@ -4075,6 +4203,10 @@
'title': 'Example',
'type': 'boolean',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'human',
'enum': list([
@ -4121,6 +4253,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'type': dict({
'default': 'system',
'enum': list([
@ -4164,6 +4300,10 @@
]),
'title': 'Content',
}),
'name': dict({
'title': 'Name',
'type': 'string',
}),
'tool_call_id': dict({
'title': 'Tool Call Id',
'type': 'string',

View File

@ -116,9 +116,9 @@ class FakeTracer(BaseTracer):
return run.copy(
update={
"id": self._replace_uuid(run.id),
"parent_run_id": self.uuids_map[run.parent_run_id]
if run.parent_run_id
else None,
"parent_run_id": (
self.uuids_map[run.parent_run_id] if run.parent_run_id else None
),
"child_runs": [self._copy_run(child) for child in run.child_runs],
"execution_order": None,
"child_execution_order": None,
@ -345,6 +345,7 @@ def test_schemas(snapshot: SnapshotAssertion) -> None:
"enum": ["ai"],
"type": "string",
},
"name": {"title": "Name", "type": "string"},
"example": {
"title": "Example",
"default": False,
@ -380,6 +381,7 @@ def test_schemas(snapshot: SnapshotAssertion) -> None:
"enum": ["human"],
"type": "string",
},
"name": {"title": "Name", "type": "string"},
"example": {
"title": "Example",
"default": False,
@ -390,7 +392,7 @@ def test_schemas(snapshot: SnapshotAssertion) -> None:
},
"ChatMessage": {
"title": "ChatMessage",
"description": "Message that can be assigned an arbitrary speaker (i.e. role).", # noqa
"description": "Message that can be assigned an arbitrary speaker (i.e. role).", # noqa: E501
"type": "object",
"properties": {
"content": {
@ -415,13 +417,14 @@ def test_schemas(snapshot: SnapshotAssertion) -> None:
"enum": ["chat"],
"type": "string",
},
"name": {"title": "Name", "type": "string"},
"role": {"title": "Role", "type": "string"},
},
"required": ["content", "role"],
},
"SystemMessage": {
"title": "SystemMessage",
"description": "Message for priming AI behavior, usually passed in as the first of a sequence\nof input messages.", # noqa
"description": "Message for priming AI behavior, usually passed in as the first of a sequence\nof input messages.", # noqa: E501
"type": "object",
"properties": {
"content": {
@ -446,12 +449,13 @@ def test_schemas(snapshot: SnapshotAssertion) -> None:
"enum": ["system"],
"type": "string",
},
"name": {"title": "Name", "type": "string"},
},
"required": ["content"],
},
"FunctionMessage": {
"title": "FunctionMessage",
"description": "Message for passing the result of executing a function back to a model.", # noqa
"description": "Message for passing the result of executing a function back to a model.", # noqa: E501
"type": "object",
"properties": {
"content": {
@ -482,7 +486,7 @@ def test_schemas(snapshot: SnapshotAssertion) -> None:
},
"ToolMessage": {
"title": "ToolMessage",
"description": "Message for passing the result of executing a tool back to a model.", # noqa
"description": "Message for passing the result of executing a tool back to a model.", # noqa: E501
"type": "object",
"properties": {
"content": {
@ -507,6 +511,7 @@ def test_schemas(snapshot: SnapshotAssertion) -> None:
"enum": ["tool"],
"type": "string",
},
"name": {"title": "Name", "type": "string"},
"tool_call_id": {"title": "Tool Call Id", "type": "string"},
},
"required": ["content", "tool_call_id"],
@ -658,14 +663,11 @@ def test_lambda_schemas() -> None:
}
second_lambda = lambda x, y: (x["hello"], x["bye"], y["bah"]) # noqa: E731
assert (
RunnableLambda(second_lambda).input_schema.schema() # type: ignore[arg-type]
== {
"title": "RunnableLambdaInput",
"type": "object",
"properties": {"hello": {"title": "Hello"}, "bye": {"title": "Bye"}},
}
)
assert RunnableLambda(second_lambda).input_schema.schema() == { # type: ignore[arg-type]
"title": "RunnableLambdaInput",
"type": "object",
"properties": {"hello": {"title": "Hello"}, "bye": {"title": "Bye"}},
}
def get_value(input): # type: ignore[no-untyped-def]
return input["variable_name"]
@ -721,7 +723,9 @@ def test_lambda_schemas() -> None:
}
assert (
RunnableLambda(aget_values_typed).input_schema.schema() # type: ignore[arg-type]
RunnableLambda(
aget_values_typed # type: ignore[arg-type]
).input_schema.schema()
== {
"title": "aget_values_typed_input",
"$ref": "#/definitions/InputType",

View File

@ -1,5 +1,5 @@
import unittest
from typing import List
from typing import List, Type
import pytest
@ -192,6 +192,21 @@ def test_multiple_msg() -> None:
assert messages_from_dict(messages_to_dict(msgs)) == msgs
def test_multiple_msg_with_name() -> None:
human_msg = HumanMessage(
content="human", additional_kwargs={"key": "value"}, name="human erick"
)
ai_msg = AIMessage(content="ai", name="ai erick")
sys_msg = SystemMessage(content="sys", name="sys erick")
msgs = [
human_msg,
ai_msg,
sys_msg,
]
assert messages_from_dict(messages_to_dict(msgs)) == msgs
def test_message_chunk_to_message() -> None:
assert message_chunk_to_message(
AIMessageChunk(content="I am", additional_kwargs={"foo": "bar"})
@ -480,3 +495,49 @@ def test_convert_to_messages() -> None:
HumanMessage(content="Hello!"),
AIMessage(content="Hi!"),
]
@pytest.mark.parametrize(
"MessageClass",
[
AIMessage,
AIMessageChunk,
HumanMessage,
HumanMessageChunk,
SystemMessage,
],
)
def test_message_name(MessageClass: Type) -> None:
msg = MessageClass(content="foo", name="bar")
assert msg.name == "bar"
msg2 = MessageClass(content="foo", name=None)
assert msg2.name is None
msg3 = MessageClass(content="foo")
assert msg3.name is None
@pytest.mark.parametrize(
"MessageClass",
[FunctionMessage, FunctionMessageChunk],
)
def test_message_name_function(MessageClass: Type) -> None:
# functionmessage doesn't support name=None
msg = MessageClass(name="foo", content="bar")
assert msg.name == "foo"
@pytest.mark.parametrize(
"MessageClass",
[ChatMessage, ChatMessageChunk],
)
def test_message_name_chat(MessageClass: Type) -> None:
msg = MessageClass(content="foo", role="user", name="bar")
assert msg.name == "bar"
msg2 = MessageClass(content="foo", role="user", name=None)
assert msg2.name is None
msg3 = MessageClass(content="foo", role="user")
assert msg3.name is None