mirror of
https://github.com/hwchase17/langchain
synced 2024-11-10 01:10:59 +00:00
4cb5f4c353
This PR should make it easier for linters to do type checking and for IDEs to jump to definition of code. See #20050 as a template for this PR. - As a byproduct: Added 3 missed `test_imports`. - Added missed `SolarChat` in to __init___.py Added it into test_import ut. - Added `# type: ignore` to fix linting. It is not clear, why linting errors appear after ^ changes. --------- Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
"""Test Anthropic Chat API wrapper."""
|
|
from typing import List
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
BaseMessage,
|
|
HumanMessage,
|
|
SystemMessage,
|
|
)
|
|
|
|
from langchain_community.chat_models import BedrockChat
|
|
from langchain_community.chat_models.meta import convert_messages_to_prompt_llama
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("messages", "expected"),
|
|
[
|
|
([HumanMessage(content="Hello")], "[INST] Hello [/INST]"),
|
|
(
|
|
[HumanMessage(content="Hello"), AIMessage(content="Answer:")],
|
|
"[INST] Hello [/INST]\nAnswer:",
|
|
),
|
|
(
|
|
[
|
|
SystemMessage(content="You're an assistant"),
|
|
HumanMessage(content="Hello"),
|
|
AIMessage(content="Answer:"),
|
|
],
|
|
"<<SYS>> You're an assistant <</SYS>>\n[INST] Hello [/INST]\nAnswer:",
|
|
),
|
|
],
|
|
)
|
|
def test_formatting(messages: List[BaseMessage], expected: str) -> None:
|
|
result = convert_messages_to_prompt_llama(messages)
|
|
assert result == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_id",
|
|
["anthropic.claude-v2", "amazon.titan-text-express-v1"],
|
|
)
|
|
def test_different_models_bedrock(model_id: str) -> None:
|
|
provider = model_id.split(".")[0]
|
|
client = MagicMock()
|
|
respbody = MagicMock()
|
|
if provider == "anthropic":
|
|
respbody.read.return_value = MagicMock(
|
|
decode=MagicMock(return_value=b'{"completion":"Hi back"}'),
|
|
)
|
|
client.invoke_model.return_value = {"body": respbody}
|
|
elif provider == "amazon":
|
|
respbody.read.return_value = '{"results": [{"outputText": "Hi back"}]}'
|
|
client.invoke_model.return_value = {"body": respbody}
|
|
|
|
model = BedrockChat(model_id=model_id, client=client)
|
|
|
|
# should not throw an error
|
|
model.invoke("hello there")
|
|
|
|
|
|
def test_bedrock_combine_llm_output() -> None:
|
|
model_id = "anthropic.claude-3-haiku-20240307-v1:0"
|
|
client = MagicMock()
|
|
llm_outputs = [
|
|
{
|
|
"model_id": "anthropic.claude-3-haiku-20240307-v1:0",
|
|
"usage": {
|
|
"completion_tokens": 1,
|
|
"prompt_tokens": 2,
|
|
"total_tokens": 3,
|
|
},
|
|
},
|
|
{
|
|
"model_id": "anthropic.claude-3-haiku-20240307-v1:0",
|
|
"usage": {
|
|
"completion_tokens": 1,
|
|
"prompt_tokens": 2,
|
|
"total_tokens": 3,
|
|
},
|
|
},
|
|
]
|
|
model = BedrockChat(model_id=model_id, client=client)
|
|
final_output = model._combine_llm_outputs(llm_outputs) # type: ignore[arg-type]
|
|
assert final_output["model_id"] == model_id
|
|
assert final_output["usage"]["completion_tokens"] == 2
|
|
assert final_output["usage"]["prompt_tokens"] == 4
|
|
assert final_output["usage"]["total_tokens"] == 6
|