#
# Copyright (c) 2024-2025 Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#

import asyncio
import json
import unittest
from typing import Any, Optional

from pipecat.audio.interruptions.min_words_interruption_strategy import MinWordsInterruptionStrategy
from pipecat.audio.turn.smart_turn.base_smart_turn import SmartTurnParams
from pipecat.audio.vad.vad_analyzer import VADParams
from pipecat.frames.frames import (
    BotStartedSpeakingFrame,
    EmulateUserStartedSpeakingFrame,
    EmulateUserStoppedSpeakingFrame,
    Frame,
    FunctionCallInProgressFrame,
    FunctionCallResultFrame,
    FunctionCallResultProperties,
    InterimTranscriptionFrame,
    InterruptionFrame,
    InterruptionTaskFrame,
    LLMContextAssistantTimestampFrame,
    LLMContextFrame,
    LLMFullResponseEndFrame,
    LLMFullResponseStartFrame,
    OpenAILLMContextAssistantTimestampFrame,
    SpeechControlParamsFrame,
    TextFrame,
    TranscriptionFrame,
    TTSTextFrame,
    UserStartedSpeakingFrame,
    UserStoppedSpeakingFrame,
)
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.task import PipelineParams
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.aggregators.llm_response import (
    LLMAssistantAggregatorParams,
    LLMAssistantContextAggregator,
    LLMUserAggregatorParams,
    LLMUserContextAggregator,
)
from pipecat.processors.aggregators.llm_response_universal import (
    LLMAssistantAggregator,
    LLMUserAggregator,
)
from pipecat.processors.aggregators.openai_llm_context import (
    OpenAILLMContext,
    OpenAILLMContextFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.services.anthropic.llm import (
    AnthropicAssistantContextAggregator,
    AnthropicLLMContext,
    AnthropicUserContextAggregator,
)
from pipecat.services.aws.llm import (
    AWSBedrockAssistantContextAggregator,
    AWSBedrockLLMContext,
    AWSBedrockUserContextAggregator,
)
from pipecat.services.google.llm import (
    GoogleAssistantContextAggregator,
    GoogleLLMContext,
    GoogleUserContextAggregator,
)
from pipecat.services.openai.llm import (
    OpenAIAssistantContextAggregator,
    OpenAIUserContextAggregator,
)
from pipecat.tests.utils import SleepFrame, run_test

AGGREGATION_TIMEOUT = 0.1
AGGREGATION_SLEEP = 0.15


class BaseTestUserContextAggregator:
    CONTEXT_CLASS = None  # To be set in subclasses
    CONTEXT_FRAME_CLASS = None  # To be set in subclasses
    USER_AGGREGATOR_CLASS = None  # To be set in subclasses
    USER_EXPECTED_CONTEXT_FRAMES = None
    ASSISTANT_AGGREGATOR_CLASS = None  # To be set in subclasses
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = None  # To be set in subclasses

    def check_message_content(self, context: OpenAILLMContext, index: int, content: str):
        assert context.messages[index]["content"] == content

    def check_message_multi_content(
        self, context: OpenAILLMContext, content_index: int, index: int, content: str
    ):
        assert context.messages[index]["content"] == content

    async def test_se(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(context)
        frames_to_send = [UserStartedSpeakingFrame(), UserStoppedSpeakingFrame()]
        expected_down_frames = [UserStartedSpeakingFrame, UserStoppedSpeakingFrame]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )

    async def test_ste(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(context)
        frames_to_send = [
            UserStartedSpeakingFrame(),
            TranscriptionFrame(text="Hello!", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
            UserStoppedSpeakingFrame,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "Hello!")

    async def test_site(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(context)
        frames_to_send = [
            UserStartedSpeakingFrame(),
            InterimTranscriptionFrame(text="Hello", user_id="cat", timestamp=""),
            TranscriptionFrame(text="Hello Pipecat!", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
            UserStoppedSpeakingFrame,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "Hello Pipecat!")

    async def test_st1iest2e(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(context)
        frames_to_send = [
            UserStartedSpeakingFrame(),
            TranscriptionFrame(text="Hello Pipecat!", user_id="cat", timestamp=""),
            InterimTranscriptionFrame(text="How ", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
            UserStartedSpeakingFrame(),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
            UserStartedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
            UserStoppedSpeakingFrame,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "Hello Pipecat! How are you?")

    async def test_siet(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=AGGREGATION_TIMEOUT)
        )
        frames_to_send = [
            UserStartedSpeakingFrame(),
            InterimTranscriptionFrame(text="How ", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "How are you?")

    async def test_sieit(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=AGGREGATION_TIMEOUT)
        )
        frames_to_send = [
            UserStartedSpeakingFrame(),
            InterimTranscriptionFrame(text="How ", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
            InterimTranscriptionFrame(text="are you?", user_id="cat", timestamp=""),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "How are you?")

    async def test_set(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=AGGREGATION_TIMEOUT)
        )
        frames_to_send = [
            UserStartedSpeakingFrame(),
            UserStoppedSpeakingFrame(),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "How are you?")

    async def test_seit(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=AGGREGATION_TIMEOUT)
        )
        frames_to_send = [
            UserStartedSpeakingFrame(),
            UserStoppedSpeakingFrame(),
            InterimTranscriptionFrame(text="How ", user_id="cat", timestamp=""),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "How are you?")

    async def test_st1et2(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=AGGREGATION_TIMEOUT)
        )
        frames_to_send = [
            SpeechControlParamsFrame(vad_params=VADParams(stop_secs=AGGREGATION_TIMEOUT)),
            UserStartedSpeakingFrame(),
            TranscriptionFrame(text="Hello Pipecat!", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            SpeechControlParamsFrame,
            UserStartedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_multi_content(context, 0, 0, "Hello Pipecat!")
        self.check_message_multi_content(context, 0, 1, "How are you?")

    async def test_set1t2(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=AGGREGATION_TIMEOUT)
        )
        frames_to_send = [
            UserStartedSpeakingFrame(),
            UserStoppedSpeakingFrame(),
            TranscriptionFrame(text="Hello Pipecat!", user_id="cat", timestamp=""),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "Hello Pipecat! How are you?")

    async def test_siet1it2(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=AGGREGATION_TIMEOUT)
        )
        frames_to_send = [
            UserStartedSpeakingFrame(),
            InterimTranscriptionFrame(text="Hello ", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
            TranscriptionFrame(text="Hello Pipecat!", user_id="cat", timestamp=""),
            InterimTranscriptionFrame(text="How ", user_id="cat", timestamp=""),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "Hello Pipecat! How are you?")

    async def test_t(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context
        )  # No aggregation timeout; this tests VAD emulation

        frames_to_send = [
            SpeechControlParamsFrame(vad_params=VADParams(stop_secs=AGGREGATION_TIMEOUT)),
            TranscriptionFrame(text="Hello!", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            SpeechControlParamsFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        expected_up_frames = [EmulateUserStartedSpeakingFrame, EmulateUserStoppedSpeakingFrame]

        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
            expected_up_frames=expected_up_frames,
        )
        self.check_message_content(context, 0, "Hello!")

    async def test_t_with_turn_analyzer(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(turn_emulated_vad_timeout=AGGREGATION_TIMEOUT)
        )

        frames_to_send = [
            SpeechControlParamsFrame(
                vad_params=VADParams(stop_secs=0.2),
                turn_params=SmartTurnParams(stop_secs=3.0),  # Turn analyzer present
            ),
            TranscriptionFrame(text="Hello!", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            SpeechControlParamsFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        expected_up_frames = [EmulateUserStartedSpeakingFrame, EmulateUserStoppedSpeakingFrame]

        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
            expected_up_frames=expected_up_frames,
        )
        self.check_message_content(context, 0, "Hello!")

    async def test_it(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context
        )  # No aggregation timeout; this tests VAD emulation
        frames_to_send = [
            SpeechControlParamsFrame(vad_params=VADParams(stop_secs=AGGREGATION_TIMEOUT)),
            InterimTranscriptionFrame(text="Hello ", user_id="cat", timestamp=""),
            SleepFrame(),
            TranscriptionFrame(text="Hello Pipecat!", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [SpeechControlParamsFrame, *self.USER_EXPECTED_CONTEXT_FRAMES]
        expected_up_frames = [EmulateUserStartedSpeakingFrame, EmulateUserStoppedSpeakingFrame]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
            expected_up_frames=expected_up_frames,
        )
        self.check_message_content(context, 0, "Hello Pipecat!")

    async def test_sie_delay_it(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=AGGREGATION_TIMEOUT)
        )
        frames_to_send = [
            UserStartedSpeakingFrame(),
            InterimTranscriptionFrame(text="How ", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
            SleepFrame(AGGREGATION_SLEEP),
            InterimTranscriptionFrame(text="are you?", user_id="cat", timestamp=""),
            TranscriptionFrame(text="How are you?", user_id="cat", timestamp=""),
            SleepFrame(sleep=AGGREGATION_SLEEP),
        ]
        expected_down_frames = [
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "How are you?")

    async def test_min_words_interruption_strategy_one_word(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.CONTEXT_FRAME_CLASS is not None, "CONTEXT_FRAME_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        CONTEXT_FRAME_CLASS = self.CONTEXT_FRAME_CLASS

        class ContextProcessor(FrameProcessor):
            def __init__(self):
                super().__init__()
                self.context_received = False

            async def process_frame(self, frame: Frame, direction: FrameDirection):
                await super().process_frame(frame, direction)

                if isinstance(frame, CONTEXT_FRAME_CLASS):
                    self.context_received = True

                await self.push_frame(frame, direction)

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(context)
        context_processor = ContextProcessor()
        pipeline = Pipeline([aggregator, context_processor])

        frames_to_send = [
            BotStartedSpeakingFrame(),
            UserStartedSpeakingFrame(),
            TranscriptionFrame(text="Can", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
        ]
        expected_down_frames = [
            BotStartedSpeakingFrame,
            UserStartedSpeakingFrame,
            UserStoppedSpeakingFrame,
        ]
        await run_test(
            pipeline,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
            pipeline_params=PipelineParams(
                interruption_strategies=[MinWordsInterruptionStrategy(min_words=2)]
            ),
        )
        assert not context_processor.context_received

    async def test_min_words_interruption_strategy_two_words(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.CONTEXT_FRAME_CLASS is not None, "CONTEXT_FRAME_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )

        CONTEXT_FRAME_CLASS = self.CONTEXT_FRAME_CLASS

        class ContextProcessor(FrameProcessor):
            def __init__(self):
                super().__init__()
                self.context_received = False

            async def process_frame(self, frame: Frame, direction: FrameDirection):
                await super().process_frame(frame, direction)

                if isinstance(frame, CONTEXT_FRAME_CLASS):
                    self.context_received = True
                elif isinstance(frame, InterruptionFrame):
                    self.context_received = False

                await self.push_frame(frame, direction)

        context = self.CONTEXT_CLASS()
        aggregator = self.USER_AGGREGATOR_CLASS(context)
        context_processor = ContextProcessor()
        pipeline = Pipeline([aggregator, context_processor])

        frames_to_send = [
            BotStartedSpeakingFrame(),
            UserStartedSpeakingFrame(),
            TranscriptionFrame(text="Can you", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
        ]
        expected_up_frames = [InterruptionTaskFrame]
        expected_down_frames = [
            BotStartedSpeakingFrame,
            UserStartedSpeakingFrame,
            InterruptionFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            pipeline,
            frames_to_send=frames_to_send,
            expected_up_frames=expected_up_frames,
            expected_down_frames=expected_down_frames,
            pipeline_params=PipelineParams(
                interruption_strategies=[MinWordsInterruptionStrategy(min_words=2)]
            ),
        )
        self.check_message_content(context, 0, "Can you")
        # If the context is not received or it has been cleared by the
        # interruption then we have an issue.
        assert context_processor.context_received

    async def test_interruption_strategy_context_order(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.USER_AGGREGATOR_CLASS is not None, (
            "USER_AGGREGATOR_CLASS must be set in a subclass"
        )
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        class DelayedProcessor(FrameProcessor):
            """Force a delay in interruption frames.

            This might give time to the assistant aggregator to update the
            context before the user aggregator (which shouldn't really happen)
            and reveal any issues in context ordering.
            """

            def __init__(self):
                super().__init__()

            async def process_frame(self, frame: Frame, direction: FrameDirection):
                await super().process_frame(frame, direction)
                if isinstance(frame, InterruptionFrame):
                    await asyncio.sleep(0.3)
                await self.push_frame(frame, direction)

        context = self.CONTEXT_CLASS()
        user_aggregator = self.USER_AGGREGATOR_CLASS(
            context, params=LLMUserAggregatorParams(aggregation_timeout=1.0)
        )
        assistant_aggregator = self.ASSISTANT_AGGREGATOR_CLASS(context)
        pipeline = Pipeline([user_aggregator, DelayedProcessor(), assistant_aggregator])

        frames_to_send = [
            # Aggregate assistant content.
            BotStartedSpeakingFrame(),
            LLMFullResponseStartFrame(),
            TTSTextFrame(text="Hello, I'm your assistant"),
            SleepFrame(),
            # Interrupt the bot. Assistant content should be added first to the
            # context, followed by user content.
            UserStartedSpeakingFrame(),
            TranscriptionFrame(text="Can you tell me", user_id="cat", timestamp=""),
            SleepFrame(),
            UserStoppedSpeakingFrame(),
        ]
        expected_down_frames = [
            BotStartedSpeakingFrame,
            UserStartedSpeakingFrame,
            *self.ASSISTANT_EXPECTED_CONTEXT_FRAMES,
            InterruptionFrame,
            UserStoppedSpeakingFrame,
            *self.USER_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            pipeline,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
            pipeline_params=PipelineParams(
                interruption_strategies=[MinWordsInterruptionStrategy(min_words=2)]
            ),
        )
        self.check_message_content(context, -1, "Can you tell me")
        self.check_message_content(context, -2, "Hello, I'm your assistant")


class BaseTestAssistantContextAggregator:
    CONTEXT_CLASS = None  # To be set in subclasses
    USER_AGGREGATOR_CLASS = None  # To be set in subclasses
    ASSISTANT_AGGREGATOR_CLASS = None  # To be set in subclasses
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = []  # To be set in subclasses

    def create_assistant_aggregator_params(
        self, **kwargs
    ) -> Optional[LLMAssistantAggregatorParams]:
        return LLMAssistantAggregatorParams(**kwargs)

    def check_message_content(self, context: OpenAILLMContext, index: int, content: str):
        assert context.messages[index]["content"] == content

    def check_message_multi_content(
        self, context: OpenAILLMContext, content_index: int, index: int, content: str
    ):
        assert context.messages[index]["content"] == content

    def check_function_call_result(self, context: OpenAILLMContext, index: int, content: str):
        assert json.loads(context.messages[index]["content"]) == content

    async def test_empty(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.ASSISTANT_AGGREGATOR_CLASS(context)
        frames_to_send = [LLMFullResponseStartFrame(), LLMFullResponseEndFrame()]
        expected_down_frames = []
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )

    async def test_single_text(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.ASSISTANT_AGGREGATOR_CLASS(context)
        frames_to_send = [
            LLMFullResponseStartFrame(),
            TextFrame(text="Hello Pipecat!"),
            LLMFullResponseEndFrame(),
        ]
        expected_down_frames = [*self.ASSISTANT_EXPECTED_CONTEXT_FRAMES]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "Hello Pipecat!")

    async def test_multiple_text(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.ASSISTANT_AGGREGATOR_CLASS(
            context, params=self.create_assistant_aggregator_params(expect_stripped_words=False)
        )
        frames_to_send = [
            LLMFullResponseStartFrame(),
            TextFrame(text="Hello "),
            TextFrame(text="Pipecat. "),
            TextFrame(text="How are "),
            TextFrame(text="you?"),
            LLMFullResponseEndFrame(),
        ]
        expected_down_frames = [*self.ASSISTANT_EXPECTED_CONTEXT_FRAMES]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "Hello Pipecat. How are you?")

    async def test_multiple_text_stripped(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.ASSISTANT_AGGREGATOR_CLASS(context)
        frames_to_send = [
            LLMFullResponseStartFrame(),
            TextFrame(text="Hello"),
            TextFrame(text="Pipecat."),
            TextFrame(text="How are"),
            TextFrame(text="you?"),
            LLMFullResponseEndFrame(),
        ]
        expected_down_frames = [*self.ASSISTANT_EXPECTED_CONTEXT_FRAMES]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_content(context, 0, "Hello Pipecat. How are you?")

    async def test_multiple_llm_responses(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.ASSISTANT_AGGREGATOR_CLASS(
            context, params=self.create_assistant_aggregator_params(expect_stripped_words=False)
        )
        frames_to_send = [
            LLMFullResponseStartFrame(),
            TextFrame(text="Hello "),
            TextFrame(text="Pipecat."),
            LLMFullResponseEndFrame(),
            LLMFullResponseStartFrame(),
            TextFrame(text="How are "),
            TextFrame(text="you?"),
            LLMFullResponseEndFrame(),
        ]
        expected_down_frames = [
            *self.ASSISTANT_EXPECTED_CONTEXT_FRAMES,
            *self.ASSISTANT_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_multi_content(context, 0, 0, "Hello Pipecat.")
        self.check_message_multi_content(context, 0, 1, "How are you?")

    async def test_multiple_llm_responses_interruption(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.ASSISTANT_AGGREGATOR_CLASS(
            context, params=self.create_assistant_aggregator_params(expect_stripped_words=False)
        )
        frames_to_send = [
            LLMFullResponseStartFrame(),
            TextFrame(text="Hello "),
            TextFrame(text="Pipecat."),
            LLMFullResponseEndFrame(),
            SleepFrame(AGGREGATION_SLEEP),
            InterruptionFrame(),
            LLMFullResponseStartFrame(),
            TextFrame(text="How are "),
            TextFrame(text="you?"),
            LLMFullResponseEndFrame(),
        ]
        expected_down_frames = [
            *self.ASSISTANT_EXPECTED_CONTEXT_FRAMES,
            InterruptionFrame,
            *self.ASSISTANT_EXPECTED_CONTEXT_FRAMES,
        ]
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_message_multi_content(context, 0, 0, "Hello Pipecat.")
        self.check_message_multi_content(context, 0, 1, "How are you?")

    async def test_function_call(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        context = self.CONTEXT_CLASS()
        aggregator = self.ASSISTANT_AGGREGATOR_CLASS(context)
        frames_to_send = [
            FunctionCallInProgressFrame(
                function_name="get_weather",
                tool_call_id="1",
                arguments={"location": "Los Angeles"},
                cancel_on_interruption=False,
            ),
            SleepFrame(),
            FunctionCallResultFrame(
                function_name="get_weather",
                tool_call_id="1",
                arguments={"location": "Los Angeles"},
                result={"conditions": "Sunny"},
            ),
        ]
        expected_down_frames = []
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_function_call_result(context, -1, {"conditions": "Sunny"})

    async def test_function_call_on_context_updated(self):
        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
        assert self.ASSISTANT_AGGREGATOR_CLASS is not None, (
            "ASSISTANT_AGGREGATOR_CLASS must be set in a subclass"
        )

        context_updated = False

        async def on_context_updated():
            nonlocal context_updated
            context_updated = True

        context = self.CONTEXT_CLASS()
        aggregator = self.ASSISTANT_AGGREGATOR_CLASS(context)
        frames_to_send = [
            FunctionCallInProgressFrame(
                function_name="get_weather",
                tool_call_id="1",
                arguments={"location": "Los Angeles"},
                cancel_on_interruption=False,
            ),
            SleepFrame(),
            FunctionCallResultFrame(
                function_name="get_weather",
                tool_call_id="1",
                arguments={"location": "Los Angeles"},
                result={"conditions": "Sunny"},
                properties=FunctionCallResultProperties(on_context_updated=on_context_updated),
            ),
            SleepFrame(),
        ]
        expected_down_frames = []
        await run_test(
            aggregator,
            frames_to_send=frames_to_send,
            expected_down_frames=expected_down_frames,
        )
        self.check_function_call_result(context, -1, {"conditions": "Sunny"})
        assert context_updated


#
# LLMUserContextAggregator
#


class TestLLMUserContextAggregator(BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase):
    CONTEXT_CLASS = OpenAILLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = LLMUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = LLMAssistantContextAggregator
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = [
        OpenAILLMContextFrame,
        OpenAILLMContextAssistantTimestampFrame,
    ]


#
# Anthropic
#


class TestAnthropicUserContextAggregator(
    BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = AnthropicLLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = AnthropicUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = AnthropicAssistantContextAggregator
    ASSISTANT_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


class TestAnthropicAssistantContextAggregator(
    BaseTestAssistantContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = AnthropicLLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = AnthropicUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = AnthropicAssistantContextAggregator
    ASSISTANT_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]["content"] == json.dumps(content)


#
# AWS (Bedrock)
#


class TestAWSBedrockUserContextAggregator(
    BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = AWSBedrockLLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = AWSBedrockUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = AWSBedrockAssistantContextAggregator
    ASSISTANT_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


class TestAWSBedrockAssistantContextAggregator(
    BaseTestAssistantContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = AWSBedrockLLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = AWSBedrockUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = AWSBedrockAssistantContextAggregator
    ASSISTANT_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
#


class TestGoogleUserContextAggregator(
    BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = GoogleLLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = GoogleUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = GoogleAssistantContextAggregator
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = [
        OpenAILLMContextFrame,
        OpenAILLMContextAssistantTimestampFrame,
    ]

    def check_message_content(self, context: OpenAILLMContext, index: int, content: str):
        obj = context.messages[index].to_json_dict()
        assert obj["parts"][0]["text"] == content

    def check_message_multi_content(
        self, context: OpenAILLMContext, content_index: int, index: int, content: str
    ):
        obj = context.messages[index].to_json_dict()
        assert obj["parts"][0]["text"] == content


class TestGoogleAssistantContextAggregator(
    BaseTestAssistantContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = GoogleLLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = GoogleUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = GoogleAssistantContextAggregator
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = [
        OpenAILLMContextFrame,
        OpenAILLMContextAssistantTimestampFrame,
    ]

    def check_message_content(self, context: OpenAILLMContext, index: int, content: str):
        obj = context.messages[index].to_json_dict()
        assert obj["parts"][0]["text"] == content

    def check_message_multi_content(
        self, context: OpenAILLMContext, content_index: int, index: int, content: str
    ):
        obj = context.messages[index].to_json_dict()
        assert obj["parts"][0]["text"] == content

    def check_function_call_result(self, context: OpenAILLMContext, index: int, content: Any):
        obj = context.messages[index].to_json_dict()
        assert obj["parts"][0]["function_response"]["response"]["value"] == json.dumps(content)


#
# OpenAI
#


class TestOpenAIUserContextAggregator(
    BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = OpenAILLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = OpenAIUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = OpenAIAssistantContextAggregator
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = [
        OpenAILLMContextFrame,
        OpenAILLMContextAssistantTimestampFrame,
    ]


class TestOpenAIAssistantContextAggregator(
    BaseTestAssistantContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = OpenAILLMContext
    CONTEXT_FRAME_CLASS = OpenAILLMContextFrame
    USER_AGGREGATOR_CLASS = OpenAIUserContextAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = OpenAIAssistantContextAggregator
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = [
        OpenAILLMContextFrame,
        OpenAILLMContextAssistantTimestampFrame,
    ]


#
# Universal
#


class TestLLMUserAggregator(BaseTestUserContextAggregator, unittest.IsolatedAsyncioTestCase):
    CONTEXT_CLASS = LLMContext
    CONTEXT_FRAME_CLASS = LLMContextFrame
    USER_AGGREGATOR_CLASS = LLMUserAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [LLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = LLMAssistantAggregator
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = [LLMContextFrame, LLMContextAssistantTimestampFrame]


class TestLLMAssistantAggregator(
    BaseTestAssistantContextAggregator, unittest.IsolatedAsyncioTestCase
):
    CONTEXT_CLASS = LLMContext
    CONTEXT_FRAME_CLASS = LLMContextFrame
    USER_AGGREGATOR_CLASS = LLMUserAggregator
    USER_EXPECTED_CONTEXT_FRAMES = [LLMContextFrame]
    ASSISTANT_AGGREGATOR_CLASS = LLMAssistantAggregator
    ASSISTANT_EXPECTED_CONTEXT_FRAMES = [LLMContextFrame, LLMContextAssistantTimestampFrame]

    # Override to remove 'expect_stripped_words' parameter, which is deprecated
    # for LLMAssistantAggregator
    def create_assistant_aggregator_params(
        self, **kwargs
    ) -> Optional[LLMAssistantAggregatorParams]:
        kwargs.pop("expect_stripped_words", None)
        return LLMAssistantAggregatorParams(**kwargs) if kwargs else None