crewAI/lib/crewai/tests/test_streaming_integration.py

"""Integration tests for streaming with real LLM interactions using cassettes."""

import pytest

from crewai import Agent, Crew, Task
from crewai.flow.flow import Flow, start
from crewai.types.streaming import CrewStreamingOutput, FlowStreamingOutput


@pytest.fixture
def researcher() -> Agent:
    """Create a researcher agent for testing."""
    return Agent(
        role="Research Analyst",
        goal="Gather comprehensive information on topics",
        backstory="You are an experienced researcher with excellent analytical skills.",
        allow_delegation=False,
    )


@pytest.fixture
def simple_task(researcher: Agent) -> Task:
    """Create a simple research task."""
    return Task(
        description="Research the latest developments in {topic}",
        expected_output="A brief summary of recent developments",
        agent=researcher,
    )


class TestStreamingCrewIntegration:
    """Integration tests for crew streaming that match documentation examples."""

    @pytest.mark.vcr()
    def test_basic_crew_streaming_from_docs(
        self, researcher: Agent, simple_task: Task
    ) -> None:
        """Test basic streaming example from documentation."""
        crew = Crew(
            agents=[researcher],
            tasks=[simple_task],
            stream=True,
            verbose=False,
        )

        streaming = crew.kickoff(inputs={"topic": "artificial intelligence"})

        assert isinstance(streaming, CrewStreamingOutput)

        chunks = []
        for chunk in streaming:
            chunks.append(chunk.content)

        assert len(chunks) > 0

        result = streaming.result
        assert result.raw is not None
        assert len(result.raw) > 0

    @pytest.mark.vcr()
    def test_streaming_with_chunk_context_from_docs(
        self, researcher: Agent, simple_task: Task
    ) -> None:
        """Test streaming with chunk context example from documentation."""
        crew = Crew(
            agents=[researcher],
            tasks=[simple_task],
            stream=True,
            verbose=False,
        )

        streaming = crew.kickoff(inputs={"topic": "AI"})

        chunk_contexts = []
        for chunk in streaming:
            chunk_contexts.append(
                {
                    "task_name": chunk.task_name,
                    "task_index": chunk.task_index,
                    "agent_role": chunk.agent_role,
                    "content": chunk.content,
                    "type": chunk.chunk_type,
                }
            )

        assert len(chunk_contexts) > 0
        assert all("agent_role" in ctx for ctx in chunk_contexts)

        result = streaming.result
        assert result is not None

    @pytest.mark.vcr()
    def test_streaming_properties_from_docs(
        self, researcher: Agent, simple_task: Task
    ) -> None:
        """Test streaming properties example from documentation."""
        crew = Crew(
            agents=[researcher],
            tasks=[simple_task],
            stream=True,
            verbose=False,
        )

        streaming = crew.kickoff(inputs={"topic": "AI"})

        for _ in streaming:
            pass

        assert streaming.is_completed is True
        full_text = streaming.get_full_text()
        assert len(full_text) > 0
        assert len(streaming.chunks) > 0

        result = streaming.result
        assert result.raw is not None

    @pytest.mark.vcr()
    @pytest.mark.asyncio
    async def test_async_streaming_from_docs(
        self, researcher: Agent, simple_task: Task
    ) -> None:
        """Test async streaming example from documentation."""
        crew = Crew(
            agents=[researcher],
            tasks=[simple_task],
            stream=True,
            verbose=False,
        )

        streaming = await crew.kickoff_async(inputs={"topic": "AI"})

        assert isinstance(streaming, CrewStreamingOutput)

        chunks = []
        async for chunk in streaming:
            chunks.append(chunk.content)

        assert len(chunks) > 0

        result = streaming.result
        assert result.raw is not None

    @pytest.mark.vcr()
    def test_kickoff_for_each_streaming_from_docs(
        self, researcher: Agent, simple_task: Task
    ) -> None:
        """Test kickoff_for_each streaming example from documentation."""
        crew = Crew(
            agents=[researcher],
            tasks=[simple_task],
            stream=True,
            verbose=False,
        )

        inputs_list = [{"topic": "AI in healthcare"}, {"topic": "AI in finance"}]

        streaming_outputs = crew.kickoff_for_each(inputs=inputs_list)

        assert len(streaming_outputs) == 2
        assert all(isinstance(s, CrewStreamingOutput) for s in streaming_outputs)

        results = []
        for streaming in streaming_outputs:
            for _ in streaming:
                pass

            result = streaming.result
            results.append(result)

        assert len(results) == 2
        assert all(r.raw is not None for r in results)


class TestStreamingFlowIntegration:
    """Integration tests for flow streaming that match documentation examples."""

    @pytest.mark.vcr()
    def test_basic_flow_streaming_from_docs(self) -> None:
        """Test basic flow streaming example from documentation."""

        class ResearchFlow(Flow):
            stream = True

            @start()
            def research_topic(self) -> str:
                researcher = Agent(
                    role="Research Analyst",
                    goal="Research topics thoroughly",
                    backstory="Expert researcher with analytical skills",
                    allow_delegation=False,
                )

                task = Task(
                    description="Research AI trends and provide insights",
                    expected_output="Detailed research findings",
                    agent=researcher,
                )

                crew = Crew(
                    agents=[researcher],
                    tasks=[task],
                    stream=True,
                    verbose=False,
                )

                streaming = crew.kickoff()
                for _ in streaming:
                    pass
                return streaming.result.raw

        flow = ResearchFlow()

        streaming = flow.kickoff()

        assert isinstance(streaming, FlowStreamingOutput)

        chunks = []
        for chunk in streaming:
            chunks.append(chunk.content)

        assert len(chunks) > 0

        result = streaming.result
        assert result is not None

    @pytest.mark.vcr()
    def test_flow_streaming_properties_from_docs(self) -> None:
        """Test flow streaming properties example from documentation."""

        class SimpleFlow(Flow):
            stream = True

            @start()
            def execute(self) -> str:
                return "Flow result"

        flow = SimpleFlow()
        streaming = flow.kickoff()

        for _ in streaming:
            pass

        assert streaming.is_completed is True
        streaming.get_full_text()
        assert len(streaming.chunks) >= 0

        result = streaming.result
        assert result is not None

    @pytest.mark.asyncio
    @pytest.mark.timeout(180)
    @pytest.mark.vcr()
    async def test_async_flow_streaming_from_docs(self) -> None:
        """Test async flow streaming example from documentation."""

        class AsyncResearchFlow(Flow):
            stream = True

            @start()
            def research(self) -> str:
                researcher = Agent(
                    role="Researcher",
                    goal="Research topics",
                    backstory="Expert researcher",
                    allow_delegation=False,
                )

                task = Task(
                    description="Research AI",
                    expected_output="Research findings",
                    agent=researcher,
                )

                crew = Crew(agents=[researcher], tasks=[task], stream=True, verbose=False)
                streaming = crew.kickoff()
                for _ in streaming:
                    pass
                return streaming.result.raw

        flow = AsyncResearchFlow()

        streaming = await flow.kickoff_async()

        assert isinstance(streaming, FlowStreamingOutput)

        chunks = []
        async for chunk in streaming:
            chunks.append(chunk.content)

        result = streaming.result
        assert result is not None