diff --git a/test-requirements.txt b/test-requirements.txt index b34a53ab9..fec8adf52 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -1 +1 @@ --e ".[anthropic,google,langchain]" +-e ".[anthropic,aws,google,langchain]" diff --git a/tests/test_context_aggregators.py b/tests/test_context_aggregators.py index dfe210e07..cd84d476d 100644 --- a/tests/test_context_aggregators.py +++ b/tests/test_context_aggregators.py @@ -40,6 +40,11 @@ from pipecat.services.anthropic.llm import ( AnthropicLLMContext, AnthropicUserContextAggregator, ) +from pipecat.services.aws.llm import ( + BedrockAssistantContextAggregator, + BedrockLLMContext, + BedrockUserContextAggregator, +) from pipecat.services.google.llm import ( GoogleAssistantContextAggregator, GoogleLLMContext, @@ -669,26 +674,6 @@ class TestLLMUserContextAggregator(BaseTestUserContextAggregator, unittest.Isola AGGREGATOR_CLASS = LLMUserContextAggregator -# -# OpenAI -# - - -class TestOpenAIUserContextAggregator( - BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase -): - CONTEXT_CLASS = OpenAILLMContext - AGGREGATOR_CLASS = OpenAIUserContextAggregator - - -class TestOpenAIAssistantContextAggregator( - BaseTestAssistantContextAggreagator, unittest.IsolatedAsyncioTestCase -): - CONTEXT_CLASS = OpenAILLMContext - AGGREGATOR_CLASS = OpenAIAssistantContextAggregator - EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame, OpenAILLMContextAssistantTimestampFrame] - - # # Anthropic # @@ -724,6 +709,43 @@ class TestAnthropicAssistantContextAggregator( assert context.messages[index]["content"][0]["content"] == json.dumps(content) +# +# AWS (Bedrock) +# + + +class TestBedrockUserContextAggregator( + BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase +): + CONTEXT_CLASS = BedrockLLMContext + AGGREGATOR_CLASS = BedrockUserContextAggregator + + def check_message_multi_content( + self, context: OpenAILLMContext, content_index: int, index: int, content: str + ): + messages = context.messages[content_index] + assert messages["content"][index]["text"] == content + + +class TestBedrockAssistantContextAggregator( + BaseTestAssistantContextAggreagator, unittest.IsolatedAsyncioTestCase +): + CONTEXT_CLASS = BedrockLLMContext + AGGREGATOR_CLASS = BedrockAssistantContextAggregator + EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame, OpenAILLMContextAssistantTimestampFrame] + + def check_message_multi_content( + self, context: OpenAILLMContext, content_index: int, index: int, content: str + ): + messages = context.messages[content_index] + assert messages["content"][index]["text"] == content + + def check_function_call_result(self, context: OpenAILLMContext, index: int, content: Any): + assert context.messages[index]["content"][0]["toolResult"]["content"][0][ + "text" + ] == json.dumps(content) + + # # Google # @@ -766,3 +788,23 @@ class TestGoogleAssistantContextAggregator( def check_function_call_result(self, context: OpenAILLMContext, index: int, content: Any): obj = glm.Content.to_dict(context.messages[index]) assert obj["parts"][0]["function_response"]["response"]["value"] == json.dumps(content) + + +# +# OpenAI +# + + +class TestOpenAIUserContextAggregator( + BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase +): + CONTEXT_CLASS = OpenAILLMContext + AGGREGATOR_CLASS = OpenAIUserContextAggregator + + +class TestOpenAIAssistantContextAggregator( + BaseTestAssistantContextAggreagator, unittest.IsolatedAsyncioTestCase +): + CONTEXT_CLASS = OpenAILLMContext + AGGREGATOR_CLASS = OpenAIAssistantContextAggregator + EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame, OpenAILLMContextAssistantTimestampFrame]