mirror of https://github.com/hwchase17/langchain
`baichuan_secret_key` use pydantic.types.SecretStr & Add Baichuan tests (#12031)
### Description - `baichuan_secret_key` use pydantic.types.SecretStr - Add Baichuan testspull/12041/head^2
parent
85bac75729
commit
8eb40b5fe2
@ -0,0 +1,40 @@
|
|||||||
|
from langchain.chat_models.baichuan import ChatBaichuan
|
||||||
|
from langchain.schema.messages import AIMessage, HumanMessage
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_baichuan() -> None:
|
||||||
|
chat = ChatBaichuan()
|
||||||
|
message = HumanMessage(content="Hello")
|
||||||
|
response = chat([message])
|
||||||
|
assert isinstance(response, AIMessage)
|
||||||
|
assert isinstance(response.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_baichuan_with_model() -> None:
|
||||||
|
chat = ChatBaichuan(model="Baichuan2-13B")
|
||||||
|
message = HumanMessage(content="Hello")
|
||||||
|
response = chat([message])
|
||||||
|
assert isinstance(response, AIMessage)
|
||||||
|
assert isinstance(response.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_baichuan_with_temperature() -> None:
|
||||||
|
chat = ChatBaichuan(model="Baichuan2-13B", temperature=1.0)
|
||||||
|
message = HumanMessage(content="Hello")
|
||||||
|
response = chat([message])
|
||||||
|
assert isinstance(response, AIMessage)
|
||||||
|
assert isinstance(response.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_baichuan_with_kwargs() -> None:
|
||||||
|
chat = ChatBaichuan()
|
||||||
|
message = HumanMessage(content="Hello")
|
||||||
|
response = chat([message], temperature=0.88, top_p=0.7)
|
||||||
|
assert isinstance(response, AIMessage)
|
||||||
|
assert isinstance(response.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
def test_extra_kwargs() -> None:
|
||||||
|
chat = ChatBaichuan(temperature=0.88, top_p=0.7)
|
||||||
|
assert chat.temperature == 0.88
|
||||||
|
assert chat.top_p == 0.7
|
@ -0,0 +1,99 @@
|
|||||||
|
import pytest
|
||||||
|
|
||||||
|
from langchain.chat_models.baichuan import (
|
||||||
|
_convert_delta_to_message_chunk,
|
||||||
|
_convert_dict_to_message,
|
||||||
|
_convert_message_to_dict,
|
||||||
|
_signature,
|
||||||
|
)
|
||||||
|
from langchain.pydantic_v1 import SecretStr
|
||||||
|
from langchain.schema.messages import (
|
||||||
|
AIMessage,
|
||||||
|
AIMessageChunk,
|
||||||
|
ChatMessage,
|
||||||
|
FunctionMessage,
|
||||||
|
HumanMessage,
|
||||||
|
HumanMessageChunk,
|
||||||
|
SystemMessage,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_message_to_dict_human() -> None:
|
||||||
|
message = HumanMessage(content="foo")
|
||||||
|
result = _convert_message_to_dict(message)
|
||||||
|
expected_output = {"role": "user", "content": "foo"}
|
||||||
|
assert result == expected_output
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_message_to_dict_ai() -> None:
|
||||||
|
message = AIMessage(content="foo")
|
||||||
|
result = _convert_message_to_dict(message)
|
||||||
|
expected_output = {"role": "assistant", "content": "foo"}
|
||||||
|
assert result == expected_output
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_message_to_dict_system() -> None:
|
||||||
|
message = SystemMessage(content="foo")
|
||||||
|
with pytest.raises(TypeError) as e:
|
||||||
|
_convert_message_to_dict(message)
|
||||||
|
assert "Got unknown type" in str(e)
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_message_to_dict_function() -> None:
|
||||||
|
message = FunctionMessage(name="foo", content="bar")
|
||||||
|
with pytest.raises(TypeError) as e:
|
||||||
|
_convert_message_to_dict(message)
|
||||||
|
assert "Got unknown type" in str(e)
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_dict_to_message_human() -> None:
|
||||||
|
message_dict = {"role": "user", "content": "foo"}
|
||||||
|
result = _convert_dict_to_message(message_dict)
|
||||||
|
expected_output = HumanMessage(content="foo")
|
||||||
|
assert result == expected_output
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_dict_to_message_ai() -> None:
|
||||||
|
message_dict = {"role": "assistant", "content": "foo"}
|
||||||
|
result = _convert_dict_to_message(message_dict)
|
||||||
|
expected_output = AIMessage(content="foo")
|
||||||
|
assert result == expected_output
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_dict_to_message_other_role() -> None:
|
||||||
|
message_dict = {"role": "system", "content": "foo"}
|
||||||
|
result = _convert_dict_to_message(message_dict)
|
||||||
|
expected_output = ChatMessage(role="system", content="foo")
|
||||||
|
assert result == expected_output
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_delta_to_message_assistant() -> None:
|
||||||
|
delta = {"role": "assistant", "content": "foo"}
|
||||||
|
result = _convert_delta_to_message_chunk(delta, AIMessageChunk)
|
||||||
|
expected_output = AIMessageChunk(content="foo")
|
||||||
|
assert result == expected_output
|
||||||
|
|
||||||
|
|
||||||
|
def test__convert_delta_to_message_human() -> None:
|
||||||
|
delta = {"role": "user", "content": "foo"}
|
||||||
|
result = _convert_delta_to_message_chunk(delta, HumanMessageChunk)
|
||||||
|
expected_output = HumanMessageChunk(content="foo")
|
||||||
|
assert result == expected_output
|
||||||
|
|
||||||
|
|
||||||
|
def test__signature() -> None:
|
||||||
|
secret_key = SecretStr("YOUR_SECRET_KEY")
|
||||||
|
|
||||||
|
result = _signature(
|
||||||
|
secret_key=secret_key,
|
||||||
|
payload={
|
||||||
|
"model": "Baichuan2-53B",
|
||||||
|
"messages": [{"role": "user", "content": "Hi"}],
|
||||||
|
},
|
||||||
|
timestamp=1697734335,
|
||||||
|
)
|
||||||
|
|
||||||
|
# The signature was generated by the demo provided by Baichuan.
|
||||||
|
# https://platform.baichuan-ai.com/docs/api#4
|
||||||
|
expected_output = "24a50b2db1648e25a244c67c5ab57d3f"
|
||||||
|
assert result == expected_output
|
Loading…
Reference in New Issue