tests: remove TestFrameProcessor, reimplement with run_test()

This commit is contained in:
Aleix Conchillo Flaqué
2025-03-18 15:52:37 -07:00
parent c15286b148
commit d1550d5a85
6 changed files with 37 additions and 269 deletions

View File

@@ -6,7 +6,7 @@
import asyncio import asyncio
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Dict, Sequence, Tuple from typing import Any, Awaitable, Callable, Dict, Optional, Sequence, Tuple
from pipecat.frames.frames import ( from pipecat.frames.frames import (
EndFrame, EndFrame,
@@ -80,8 +80,8 @@ async def run_test(
processor: FrameProcessor, processor: FrameProcessor,
*, *,
frames_to_send: Sequence[Frame], frames_to_send: Sequence[Frame],
expected_down_frames: Sequence[type], expected_down_frames: Optional[Sequence[type]] = None,
expected_up_frames: Sequence[type] = [], expected_up_frames: Optional[Sequence[type]] = None,
ignore_start: bool = True, ignore_start: bool = True,
start_metadata: Dict[str, Any] = {}, start_metadata: Dict[str, Any] = {},
send_end_frame: bool = True, send_end_frame: bool = True,
@@ -126,33 +126,35 @@ async def run_test(
# Down frames # Down frames
# #
received_down_frames: Sequence[Frame] = [] received_down_frames: Sequence[Frame] = []
while not received_down.empty(): if expected_down_frames is not None:
frame = await received_down.get() while not received_down.empty():
if not isinstance(frame, EndFrame) or not send_end_frame: frame = await received_down.get()
received_down_frames.append(frame) if not isinstance(frame, EndFrame) or not send_end_frame:
received_down_frames.append(frame)
print("received DOWN frames =", received_down_frames) print("received DOWN frames =", received_down_frames)
print("expected DOWN frames =", expected_down_frames) print("expected DOWN frames =", expected_down_frames)
assert len(received_down_frames) == len(expected_down_frames) assert len(received_down_frames) == len(expected_down_frames)
for real, expected in zip(received_down_frames, expected_down_frames): for real, expected in zip(received_down_frames, expected_down_frames):
assert isinstance(real, expected) assert isinstance(real, expected)
# #
# Up frames # Up frames
# #
received_up_frames: Sequence[Frame] = [] received_up_frames: Sequence[Frame] = []
while not received_up.empty(): if expected_up_frames is not None:
frame = await received_up.get() while not received_up.empty():
received_up_frames.append(frame) frame = await received_up.get()
received_up_frames.append(frame)
print("received UP frames =", received_up_frames) print("received UP frames =", received_up_frames)
print("expected UP frames =", expected_up_frames) print("expected UP frames =", expected_up_frames)
assert len(received_up_frames) == len(expected_up_frames) assert len(received_up_frames) == len(expected_up_frames)
for real, expected in zip(received_up_frames, expected_up_frames): for real, expected in zip(received_up_frames, expected_up_frames):
assert isinstance(real, expected) assert isinstance(real, expected)
return (received_down_frames, received_up_frames) return (received_down_frames, received_up_frames)

View File

@@ -1,48 +0,0 @@
from typing import List
from pipecat.processors.frame_processor import FrameProcessor
class TestException(Exception):
pass
class TestFrameProcessor(FrameProcessor):
__test__ = False # Prevents pytest from collecting this class as a test
def __init__(self, test_frames):
self.test_frames = test_frames
self._list_counter = 0
super().__init__()
async def process_frame(self, frame, direction):
await super().process_frame(frame, direction)
if not self.test_frames[
0
]: # then we've run out of required frames but the generator is still going?
raise TestException(f"Oops, got an extra frame, {frame}")
if isinstance(self.test_frames[0], List):
# We need to consume frames until we see the next frame type after this
next_frame = self.test_frames[1]
if isinstance(frame, next_frame):
# we're done iterating the list I guess
print(f"TestFrameProcessor got expected list exit frame: {frame}")
# pop twice to get rid of the list, as well as the next frame
self.test_frames.pop(0)
self.test_frames.pop(0)
self.list_counter = 0
else:
fl = self.test_frames[0]
fl_el = fl[self._list_counter % len(fl)]
if isinstance(frame, fl_el):
print(f"TestFrameProcessor got expected list frame: {frame}")
self._list_counter += 1
else:
raise TestException(f"Inside a list, expected {fl_el} but got {frame}")
else:
if not isinstance(frame, self.test_frames[0]):
raise TestException(f"Expected {self.test_frames[0]}, but got {frame}")
print(f"TestFrameProcessor got expected frame: {frame}")
self.test_frames.pop(0)

View File

@@ -1,33 +0,0 @@
import asyncio
import os
import unittest
from openai.types.chat import (
ChatCompletionSystemMessageParam,
)
from pipecat.processors.aggregators.openai_llm_context import (
OpenAILLMContext,
OpenAILLMContextFrame,
)
from pipecat.services.azure import AzureLLMService
if __name__ == "__main__":
@unittest.skip("Skip azure integration test")
async def test_chat():
llm = AzureLLMService(
api_key=os.getenv("AZURE_CHATGPT_API_KEY"),
endpoint=os.getenv("AZURE_CHATGPT_ENDPOINT"),
model=os.getenv("AZURE_CHATGPT_MODEL"),
)
context = OpenAILLMContext()
message: ChatCompletionSystemMessageParam = ChatCompletionSystemMessageParam(
content="Please tell the world hello.", name="system", role="system"
)
context.add_message(message)
frame = OpenAILLMContextFrame(context)
async for s in llm.process_frame(frame):
print(s)
asyncio.run(test_chat())

View File

@@ -1,28 +0,0 @@
import asyncio
import unittest
from openai.types.chat import (
ChatCompletionSystemMessageParam,
)
from pipecat.processors.aggregators.openai_llm_context import (
OpenAILLMContext,
OpenAILLMContextFrame,
)
from pipecat.services.ollama import OLLamaLLMService
if __name__ == "__main__":
@unittest.skip("Skip azure integration test")
async def test_chat():
llm = OLLamaLLMService()
context = OpenAILLMContext()
message: ChatCompletionSystemMessageParam = ChatCompletionSystemMessageParam(
content="Please tell the world hello.", name="system", role="system"
)
context.add_message(message)
frame = OpenAILLMContextFrame(context)
async for s in llm.process_frame(frame):
print(s)
asyncio.run(test_chat())

View File

@@ -1,128 +0,0 @@
import asyncio
import json
import os
from typing import List
from openai.types.chat import (
ChatCompletionSystemMessageParam,
ChatCompletionToolParam,
ChatCompletionUserMessageParam,
)
from pipecat.frames.frames import LLMFullResponseEndFrame, LLMFullResponseStartFrame, TextFrame
from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.openai import OpenAILLMContext, OpenAILLMContextFrame, OpenAILLMService
from pipecat.utils.test_frame_processor import TestFrameProcessor
tools = [
ChatCompletionToolParam(
type="function",
function={
"name": "get_current_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"format": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
"description": "The temperature unit to use. Infer this from the users location.",
},
},
"required": ["location", "format"],
},
},
)
]
if __name__ == "__main__":
async def test_simple_functions():
async def get_weather_from_api(llm, args):
return json.dumps({"conditions": "nice", "temperature": "75"})
api_key = os.getenv("OPENAI_API_KEY")
llm = OpenAILLMService(
api_key=api_key or "",
model="gpt-4-1106-preview",
)
llm.register_function("get_current_weather", get_weather_from_api)
t = TestFrameProcessor([LLMFullResponseStartFrame, TextFrame, LLMFullResponseEndFrame])
llm.link(t)
context = OpenAILLMContext(tools=tools)
system_message: ChatCompletionSystemMessageParam = ChatCompletionSystemMessageParam(
content="Ask the user to ask for a weather report", name="system", role="system"
)
user_message: ChatCompletionUserMessageParam = ChatCompletionUserMessageParam(
content="Could you tell me the weather for Boulder, Colorado",
name="user",
role="user",
)
context.add_message(system_message)
context.add_message(user_message)
frame = OpenAILLMContextFrame(context)
await llm.process_frame(frame, FrameDirection.DOWNSTREAM)
async def test_advanced_functions():
async def get_weather_from_api(llm, args):
return [
{
"role": "system",
"content": "The user has asked for live weather. Respond by telling them we don't currently support live weather for that area, but it's coming soon.",
}
]
api_key = os.getenv("OPENAI_API_KEY")
llm = OpenAILLMService(
api_key=api_key or "",
model="gpt-4-1106-preview",
)
llm.register_function("get_current_weather", get_weather_from_api)
t = TestFrameProcessor([LLMFullResponseStartFrame, TextFrame, LLMFullResponseEndFrame])
llm.link(t)
context = OpenAILLMContext(tools=tools)
system_message: ChatCompletionSystemMessageParam = ChatCompletionSystemMessageParam(
content="Ask the user to ask for a weather report", name="system", role="system"
)
user_message: ChatCompletionUserMessageParam = ChatCompletionUserMessageParam(
content="Could you tell me the weather for Boulder, Colorado",
name="user",
role="user",
)
context.add_message(system_message)
context.add_message(user_message)
frame = OpenAILLMContextFrame(context)
await llm.process_frame(frame, FrameDirection.DOWNSTREAM)
async def test_chat():
api_key = os.getenv("OPENAI_API_KEY")
t = TestFrameProcessor([LLMFullResponseStartFrame, TextFrame, LLMFullResponseEndFrame])
llm = OpenAILLMService(
api_key=api_key or "",
model="gpt-4o",
)
llm.link(t)
context = OpenAILLMContext()
message: ChatCompletionSystemMessageParam = ChatCompletionSystemMessageParam(
content="Please tell the world hello.", name="system", role="system"
)
context.add_message(message)
frame = OpenAILLMContextFrame(context)
await llm.process_frame(frame, FrameDirection.DOWNSTREAM)
async def run_tests():
await test_simple_functions()
await test_advanced_functions()
await test_chat()
asyncio.run(run_tests())

View File

@@ -1,3 +1,9 @@
#
# Copyright (c) 2024-2025 Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#
import os import os
from unittest.mock import AsyncMock from unittest.mock import AsyncMock
@@ -6,17 +12,12 @@ from dotenv import load_dotenv
from pipecat.adapters.schemas.function_schema import FunctionSchema from pipecat.adapters.schemas.function_schema import FunctionSchema
from pipecat.adapters.schemas.tools_schema import ToolsSchema from pipecat.adapters.schemas.tools_schema import ToolsSchema
from pipecat.frames.frames import ( from pipecat.pipeline.pipeline import Pipeline
LLMFullResponseEndFrame,
LLMFullResponseStartFrame,
LLMTextFrame,
)
from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.ai_services import LLMService from pipecat.services.ai_services import LLMService
from pipecat.services.anthropic import AnthropicLLMService from pipecat.services.anthropic import AnthropicLLMService
from pipecat.services.google import GoogleLLMService from pipecat.services.google import GoogleLLMService
from pipecat.services.openai import OpenAILLMContext, OpenAILLMContextFrame, OpenAILLMService from pipecat.services.openai import OpenAILLMContext, OpenAILLMContextFrame, OpenAILLMService
from pipecat.utils.test_frame_processor import TestFrameProcessor from pipecat.tests.utils import run_test
load_dotenv(override=True) load_dotenv(override=True)
@@ -47,8 +48,6 @@ async def _test_llm_function_calling(llm: LLMService):
mock_fetch_weather = AsyncMock() mock_fetch_weather = AsyncMock()
llm.register_function(None, mock_fetch_weather) llm.register_function(None, mock_fetch_weather)
t = TestFrameProcessor([LLMFullResponseStartFrame, LLMTextFrame, LLMFullResponseEndFrame])
llm.link(t)
messages = [ messages = [
{ {
@@ -61,10 +60,14 @@ async def _test_llm_function_calling(llm: LLMService):
# This is done by default inside the create_context_aggregator # This is done by default inside the create_context_aggregator
context.set_llm_adapter(llm.get_llm_adapter()) context.set_llm_adapter(llm.get_llm_adapter())
frame = OpenAILLMContextFrame(context) pipeline = Pipeline([llm])
# This will fail if an exception is raised frames_to_send = [OpenAILLMContextFrame(context)]
await llm.process_frame(frame, FrameDirection.DOWNSTREAM) await run_test(
pipeline,
frames_to_send=frames_to_send,
expected_down_frames=None,
)
# Assert that the mock function was called # Assert that the mock function was called
mock_fetch_weather.assert_called_once() mock_fetch_weather.assert_called_once()