mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-02 07:42:40 +00:00
feat: add streaming result support to flows and crews
* feat: add streaming result support to flows and crews * docs: add streaming execution documentation and integration tests
This commit is contained in:
Greyson LaLonde
290
lib/crewai/tests/test_streaming_integration.py
Normal file
290
lib/crewai/tests/test_streaming_integration.py
Normal file
@@ -0,0 +1,290 @@
|
||||
"""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(filter_headers=["authorization"])
|
||||
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(filter_headers=["authorization"])
|
||||
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(filter_headers=["authorization"])
|
||||
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(filter_headers=["authorization"])
|
||||
@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(filter_headers=["authorization"])
|
||||
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(filter_headers=["authorization"])
|
||||
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(filter_headers=["authorization"])
|
||||
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.vcr(filter_headers=["authorization"])
|
||||
@pytest.mark.asyncio
|
||||
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
|
||||
Reference in New Issue
Block a user