mirror of
https://github.com/hwchase17/langchain
synced 2024-11-06 03:20:49 +00:00
aad2aa7188
Titan Express model was not supported as a chat model because LangChain messages were not "translated" to a text prompt. Co-authored-by: Guillem Orellana Trullols <guillem.orellana_trullols@siemens.com>
61 lines
1.8 KiB
Python
61 lines
1.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")
|