|
|
|
@ -1,8 +1,10 @@
|
|
|
|
|
"""Test Anthropic Chat API wrapper."""
|
|
|
|
|
from typing import List
|
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
|
|
|
|
|
|
import pytest
|
|
|
|
|
|
|
|
|
|
from langchain.chat_models import BedrockChat
|
|
|
|
|
from langchain.chat_models.meta import convert_messages_to_prompt_llama
|
|
|
|
|
from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
|
|
|
|
|
|
|
|
|
@ -28,3 +30,19 @@ from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
|
|
|
|
|
def test_formatting(messages: List[BaseMessage], expected: str) -> None:
|
|
|
|
|
result = convert_messages_to_prompt_llama(messages)
|
|
|
|
|
assert result == expected
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_anthropic_bedrock() -> None:
|
|
|
|
|
client = MagicMock()
|
|
|
|
|
respbody = MagicMock(
|
|
|
|
|
read=MagicMock(
|
|
|
|
|
return_value=MagicMock(
|
|
|
|
|
decode=MagicMock(return_value=b'{"completion":"Hi back"}')
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
client.invoke_model.return_value = {"body": respbody}
|
|
|
|
|
model = BedrockChat(model_id="anthropic.claude-v2", client=client)
|
|
|
|
|
|
|
|
|
|
# should not throw an error
|
|
|
|
|
model.invoke("hello there")
|
|
|
|
|