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