mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-09 14:35:50 +00:00
Bedrock scheduled tests (#11194)
This commit is contained in:
parent
f9b565fa8c
commit
38d5b63a10
8
.github/workflows/scheduled_test.yml
vendored
8
.github/workflows/scheduled_test.yml
vendored
@ -40,6 +40,13 @@ jobs:
|
|||||||
with:
|
with:
|
||||||
credentials_json: '${{ secrets.GOOGLE_CREDENTIALS }}'
|
credentials_json: '${{ secrets.GOOGLE_CREDENTIALS }}'
|
||||||
|
|
||||||
|
- name: Configure AWS Credentials
|
||||||
|
uses: aws-actions/configure-aws-credentials@v4
|
||||||
|
with:
|
||||||
|
aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
|
||||||
|
aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
|
||||||
|
aws-region: ${{ vars.AWS_REGION }}
|
||||||
|
|
||||||
- name: Install dependencies
|
- name: Install dependencies
|
||||||
working-directory: libs/langchain
|
working-directory: libs/langchain
|
||||||
shell: bash
|
shell: bash
|
||||||
@ -47,6 +54,7 @@ jobs:
|
|||||||
echo "Running scheduled tests, installing dependencies with poetry..."
|
echo "Running scheduled tests, installing dependencies with poetry..."
|
||||||
poetry install --with=test_integration
|
poetry install --with=test_integration
|
||||||
poetry run pip install google-cloud-aiplatform
|
poetry run pip install google-cloud-aiplatform
|
||||||
|
poetry run pip install "boto3>=1.28.57"
|
||||||
|
|
||||||
- name: Run tests
|
- name: Run tests
|
||||||
shell: bash
|
shell: bash
|
||||||
|
@ -0,0 +1,143 @@
|
|||||||
|
"""Test Bedrock chat model."""
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from langchain.callbacks.manager import CallbackManager
|
||||||
|
from langchain.chat_models import BedrockChat
|
||||||
|
from langchain.schema import ChatGeneration, LLMResult
|
||||||
|
from langchain.schema.messages import BaseMessage, HumanMessage, SystemMessage
|
||||||
|
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def chat() -> BedrockChat:
|
||||||
|
return BedrockChat(model_id="anthropic.claude-v2", model_kwargs={"temperature": 0})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
def test_chat_bedrock(chat: BedrockChat) -> None:
|
||||||
|
"""Test BedrockChat wrapper."""
|
||||||
|
system = SystemMessage(content="You are a helpful assistant.")
|
||||||
|
human = HumanMessage(content="Hello")
|
||||||
|
response = chat([system, human])
|
||||||
|
assert isinstance(response, BaseMessage)
|
||||||
|
assert isinstance(response.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
def test_chat_bedrock_generate(chat: BedrockChat) -> None:
|
||||||
|
"""Test BedrockChat wrapper with generate."""
|
||||||
|
message = HumanMessage(content="Hello")
|
||||||
|
response = chat.generate([[message], [message]])
|
||||||
|
assert isinstance(response, LLMResult)
|
||||||
|
assert len(response.generations) == 2
|
||||||
|
for generations in response.generations:
|
||||||
|
for generation in generations:
|
||||||
|
assert isinstance(generation, ChatGeneration)
|
||||||
|
assert isinstance(generation.text, str)
|
||||||
|
assert generation.text == generation.message.content
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
def test_chat_bedrock_streaming() -> None:
|
||||||
|
"""Test that streaming correctly invokes on_llm_new_token callback."""
|
||||||
|
callback_handler = FakeCallbackHandler()
|
||||||
|
callback_manager = CallbackManager([callback_handler])
|
||||||
|
chat = BedrockChat(
|
||||||
|
model_id="anthropic.claude-v2",
|
||||||
|
streaming=True,
|
||||||
|
callback_manager=callback_manager,
|
||||||
|
verbose=True,
|
||||||
|
)
|
||||||
|
message = HumanMessage(content="Hello")
|
||||||
|
response = chat([message])
|
||||||
|
assert callback_handler.llm_streams > 0
|
||||||
|
assert isinstance(response, BaseMessage)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
def test_chat_bedrock_streaming_generation_info() -> None:
|
||||||
|
"""Test that generation info is preserved when streaming."""
|
||||||
|
|
||||||
|
class _FakeCallback(FakeCallbackHandler):
|
||||||
|
saved_things: dict = {}
|
||||||
|
|
||||||
|
def on_llm_end(
|
||||||
|
self,
|
||||||
|
*args: Any,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
# Save the generation
|
||||||
|
self.saved_things["generation"] = args[0]
|
||||||
|
|
||||||
|
callback = _FakeCallback()
|
||||||
|
callback_manager = CallbackManager([callback])
|
||||||
|
chat = BedrockChat(
|
||||||
|
model_id="anthropic.claude-v2",
|
||||||
|
callback_manager=callback_manager,
|
||||||
|
)
|
||||||
|
list(chat.stream("hi"))
|
||||||
|
generation = callback.saved_things["generation"]
|
||||||
|
# `Hello!` is two tokens, assert that that is what is returned
|
||||||
|
assert generation.generations[0][0].text == " Hello!"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
def test_bedrock_streaming(chat: BedrockChat) -> None:
|
||||||
|
"""Test streaming tokens from OpenAI."""
|
||||||
|
|
||||||
|
for token in chat.stream("I'm Pickle Rick"):
|
||||||
|
assert isinstance(token.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bedrock_astream(chat: BedrockChat) -> None:
|
||||||
|
"""Test streaming tokens from OpenAI."""
|
||||||
|
|
||||||
|
async for token in chat.astream("I'm Pickle Rick"):
|
||||||
|
assert isinstance(token.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bedrock_abatch(chat: BedrockChat) -> None:
|
||||||
|
"""Test streaming tokens from BedrockChat."""
|
||||||
|
result = await chat.abatch(["I'm Pickle Rick", "I'm not Pickle Rick"])
|
||||||
|
for token in result:
|
||||||
|
assert isinstance(token.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bedrock_abatch_tags(chat: BedrockChat) -> None:
|
||||||
|
"""Test batch tokens from BedrockChat."""
|
||||||
|
result = await chat.abatch(
|
||||||
|
["I'm Pickle Rick", "I'm not Pickle Rick"], config={"tags": ["foo"]}
|
||||||
|
)
|
||||||
|
for token in result:
|
||||||
|
assert isinstance(token.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
def test_bedrock_batch(chat: BedrockChat) -> None:
|
||||||
|
"""Test batch tokens from BedrockChat."""
|
||||||
|
result = chat.batch(["I'm Pickle Rick", "I'm not Pickle Rick"])
|
||||||
|
for token in result:
|
||||||
|
assert isinstance(token.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bedrock_ainvoke(chat: BedrockChat) -> None:
|
||||||
|
"""Test invoke tokens from BedrockChat."""
|
||||||
|
result = await chat.ainvoke("I'm Pickle Rick", config={"tags": ["foo"]})
|
||||||
|
assert isinstance(result.content, str)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.scheduled
|
||||||
|
def test_bedrock_invoke(chat: BedrockChat) -> None:
|
||||||
|
"""Test invoke tokens from BedrockChat."""
|
||||||
|
result = chat.invoke("I'm Pickle Rick", config=dict(tags=["foo"]))
|
||||||
|
assert isinstance(result.content, str)
|
Loading…
Reference in New Issue
Block a user