Move dedicated LLM summarization into LLMContextSummarizer
The dedicated LLM logic lived in LLMAssistantAggregator, creating two code paths and requiring the aggregator to call a private LLMService method. Move it into the summarizer which already owns the config and summarization lifecycle, keeping the aggregator handler as a single-line upstream push.
This commit is contained in:
Mark Backman
@@ -6,9 +6,10 @@
|
||||
|
||||
"""This module defines a summarizer for managing LLM context summarization."""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@@ -27,6 +28,9 @@ from pipecat.utils.context.llm_context_summarization import (
|
||||
LLMContextSummarizationUtil,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pipecat.services.llm_service import LLMService
|
||||
|
||||
|
||||
@dataclass
|
||||
class SummaryAppliedEvent:
|
||||
@@ -227,8 +231,10 @@ class LLMContextSummarizer(BaseObject):
|
||||
async def _request_summarization(self):
|
||||
"""Request context summarization from LLM service.
|
||||
|
||||
Creates a summarization request frame and emits it via event handler.
|
||||
Tracks the request ID to match async responses and prevent race conditions.
|
||||
Creates a summarization request frame and either handles it directly
|
||||
using a dedicated LLM (if configured) or emits it via event handler
|
||||
for the pipeline's primary LLM. Tracks the request ID to match async
|
||||
responses and prevent race conditions.
|
||||
"""
|
||||
# Generate unique request ID
|
||||
request_id = str(uuid.uuid4())
|
||||
@@ -250,8 +256,61 @@ class LLMContextSummarizer(BaseObject):
|
||||
summarization_timeout=self._config.summarization_timeout,
|
||||
)
|
||||
|
||||
# Emit event for aggregator to broadcast
|
||||
await self._call_event_handler("on_request_summarization", request_frame)
|
||||
if self._config.llm:
|
||||
# Use dedicated LLM directly — no need to involve the pipeline
|
||||
self.task_manager.create_task(
|
||||
self._generate_summary_with_dedicated_llm(self._config.llm, request_frame),
|
||||
f"{self}-dedicated-llm-summary",
|
||||
)
|
||||
else:
|
||||
# Emit event for aggregator to broadcast to the pipeline LLM
|
||||
await self._call_event_handler("on_request_summarization", request_frame)
|
||||
|
||||
async def _generate_summary_with_dedicated_llm(
|
||||
self, llm: "LLMService", frame: LLMContextSummaryRequestFrame
|
||||
):
|
||||
"""Generate summary using a dedicated LLM service.
|
||||
|
||||
Calls the dedicated LLM's _generate_summary directly and feeds the
|
||||
result back through _handle_summary_result, bypassing the pipeline.
|
||||
|
||||
Args:
|
||||
llm: The dedicated LLM service to use for summarization.
|
||||
frame: The summarization request frame.
|
||||
"""
|
||||
try:
|
||||
if frame.summarization_timeout:
|
||||
summary, last_index = await asyncio.wait_for(
|
||||
llm._generate_summary(frame),
|
||||
timeout=frame.summarization_timeout,
|
||||
)
|
||||
else:
|
||||
summary, last_index = await llm._generate_summary(frame)
|
||||
result_frame = LLMContextSummaryResultFrame(
|
||||
request_id=frame.request_id,
|
||||
summary=summary,
|
||||
last_summarized_index=last_index,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
error = f"Context summarization timed out after {frame.summarization_timeout}s"
|
||||
logger.error(f"{self}: {error}")
|
||||
result_frame = LLMContextSummaryResultFrame(
|
||||
request_id=frame.request_id,
|
||||
summary="",
|
||||
last_summarized_index=-1,
|
||||
error=error,
|
||||
)
|
||||
except Exception as e:
|
||||
error = f"Error generating context summary: {e}"
|
||||
logger.error(f"{self}: {error}")
|
||||
result_frame = LLMContextSummaryResultFrame(
|
||||
request_id=frame.request_id,
|
||||
summary="",
|
||||
last_summarized_index=-1,
|
||||
error=error,
|
||||
)
|
||||
|
||||
await self._handle_summary_result(result_frame)
|
||||
|
||||
async def _handle_summary_result(self, frame: LLMContextSummaryResultFrame):
|
||||
"""Handle context summarization result from LLM service.
|
||||
|
||||
@@ -16,7 +16,7 @@ import json
|
||||
import warnings
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set, Type
|
||||
from typing import Any, Dict, List, Literal, Optional, Set, Type
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@@ -39,7 +39,6 @@ from pipecat.frames.frames import (
|
||||
LLMContextAssistantTimestampFrame,
|
||||
LLMContextFrame,
|
||||
LLMContextSummaryRequestFrame,
|
||||
LLMContextSummaryResultFrame,
|
||||
LLMFullResponseEndFrame,
|
||||
LLMFullResponseStartFrame,
|
||||
LLMMessagesAppendFrame,
|
||||
@@ -84,9 +83,6 @@ from pipecat.utils.context.llm_context_summarization import LLMContextSummarizat
|
||||
from pipecat.utils.string import TextPartForConcatenation, concatenate_aggregated_text
|
||||
from pipecat.utils.time import time_now_iso8601
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pipecat.services.llm_service import LLMService
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMUserAggregatorParams:
|
||||
@@ -1252,71 +1248,13 @@ class LLMAssistantAggregator(LLMContextAggregator):
|
||||
):
|
||||
"""Handle summarization request from the summarizer.
|
||||
|
||||
If a dedicated summarization LLM is configured, generates the summary
|
||||
directly and feeds the result to the summarizer. Otherwise, pushes the
|
||||
request frame upstream to the pipeline's primary LLM service.
|
||||
Push the request frame UPSTREAM to the LLM service for processing.
|
||||
|
||||
Args:
|
||||
summarizer: The summarizer that generated the request.
|
||||
frame: The summarization request frame to broadcast.
|
||||
"""
|
||||
summarization_llm = (
|
||||
self._params.context_summarization_config.llm
|
||||
if self._params.context_summarization_config
|
||||
else None
|
||||
)
|
||||
|
||||
if summarization_llm:
|
||||
self.create_task(self._generate_summary_with_dedicated_llm(summarization_llm, frame))
|
||||
else:
|
||||
await self.push_frame(frame, FrameDirection.UPSTREAM)
|
||||
|
||||
async def _generate_summary_with_dedicated_llm(
|
||||
self, llm: "LLMService", frame: LLMContextSummaryRequestFrame
|
||||
):
|
||||
"""Generate summary using a dedicated LLM service.
|
||||
|
||||
Calls the dedicated LLM's _generate_summary directly and feeds the
|
||||
result back to the summarizer, bypassing the pipeline.
|
||||
|
||||
Args:
|
||||
llm: The dedicated LLM service to use for summarization.
|
||||
frame: The summarization request frame.
|
||||
"""
|
||||
try:
|
||||
if frame.summarization_timeout:
|
||||
summary, last_index = await asyncio.wait_for(
|
||||
llm._generate_summary(frame),
|
||||
timeout=frame.summarization_timeout,
|
||||
)
|
||||
else:
|
||||
summary, last_index = await llm._generate_summary(frame)
|
||||
result_frame = LLMContextSummaryResultFrame(
|
||||
request_id=frame.request_id,
|
||||
summary=summary,
|
||||
last_summarized_index=last_index,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
error = f"Context summarization timed out after {frame.summarization_timeout}s"
|
||||
logger.error(f"{self}: {error}")
|
||||
result_frame = LLMContextSummaryResultFrame(
|
||||
request_id=frame.request_id,
|
||||
summary="",
|
||||
last_summarized_index=-1,
|
||||
error=error,
|
||||
)
|
||||
except Exception as e:
|
||||
error = f"Error generating context summary: {e}"
|
||||
await self.push_error(error_msg=error, exception=e)
|
||||
result_frame = LLMContextSummaryResultFrame(
|
||||
request_id=frame.request_id,
|
||||
summary="",
|
||||
last_summarized_index=-1,
|
||||
error=error,
|
||||
)
|
||||
|
||||
if self._summarizer:
|
||||
await self._summarizer.process_frame(result_frame)
|
||||
await self.push_frame(frame, FrameDirection.UPSTREAM)
|
||||
|
||||
|
||||
class LLMContextAggregatorPair:
|
||||
|
||||
@@ -654,85 +654,79 @@ class TestSummaryGenerationExceptions(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
|
||||
class TestDedicatedLLMSummarization(unittest.IsolatedAsyncioTestCase):
|
||||
"""Tests for dedicated LLM summarization in LLMAssistantAggregator."""
|
||||
"""Tests for dedicated LLM summarization in LLMContextSummarizer."""
|
||||
|
||||
def _create_context_and_frame(self):
|
||||
"""Create a context with enough messages and a matching request frame."""
|
||||
context = LLMContext()
|
||||
context.add_message({"role": "user", "content": "Message 1"})
|
||||
context.add_message({"role": "assistant", "content": "Response 1"})
|
||||
context.add_message({"role": "user", "content": "Message 2"})
|
||||
|
||||
frame = LLMContextSummaryRequestFrame(
|
||||
request_id="dedicated_test",
|
||||
context=context,
|
||||
min_messages_to_keep=1,
|
||||
target_context_tokens=1000,
|
||||
summarization_prompt="Summarize this",
|
||||
summarization_timeout=5.0,
|
||||
)
|
||||
return context, frame
|
||||
|
||||
async def test_dedicated_llm_success(self):
|
||||
"""Test that dedicated LLM generates summary and feeds result to summarizer."""
|
||||
from pipecat.processors.aggregators.llm_context_summarizer import LLMContextSummarizer
|
||||
from pipecat.processors.aggregators.llm_response_universal import (
|
||||
LLMAssistantAggregator,
|
||||
LLMAssistantAggregatorParams,
|
||||
)
|
||||
async def asyncSetUp(self):
|
||||
from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams
|
||||
|
||||
context, frame = self._create_context_and_frame()
|
||||
self.task_manager = TaskManager()
|
||||
self.task_manager.setup(TaskManagerParams(loop=asyncio.get_running_loop()))
|
||||
|
||||
# Create a mock dedicated LLM
|
||||
dedicated_llm = LLMService()
|
||||
dedicated_llm._generate_summary = AsyncMock(return_value=("Dedicated summary", 1))
|
||||
def _create_context_and_config(self, dedicated_llm):
|
||||
"""Create a context with enough messages and a config with a dedicated LLM."""
|
||||
context = LLMContext()
|
||||
for i in range(10):
|
||||
context.add_message(
|
||||
{"role": "user", "content": f"Test message {i} that adds tokens to context."}
|
||||
)
|
||||
|
||||
config = LLMContextSummarizationConfig(
|
||||
max_context_tokens=50,
|
||||
max_context_tokens=50, # Very low to trigger easily
|
||||
llm=dedicated_llm,
|
||||
summarization_timeout=5.0,
|
||||
)
|
||||
params = LLMAssistantAggregatorParams(
|
||||
enable_context_summarization=True,
|
||||
context_summarization_config=config,
|
||||
)
|
||||
aggregator = LLMAssistantAggregator(context, params=params)
|
||||
return context, config
|
||||
|
||||
# Mock summarizer.process_frame to capture the result
|
||||
result_frames = []
|
||||
original_process = aggregator._summarizer.process_frame
|
||||
async def test_dedicated_llm_success(self):
|
||||
"""Test that dedicated LLM generates summary and applies result."""
|
||||
from pipecat.processors.aggregators.llm_context_summarizer import LLMContextSummarizer
|
||||
|
||||
async def capture_process(frame):
|
||||
result_frames.append(frame)
|
||||
await original_process(frame)
|
||||
dedicated_llm = LLMService()
|
||||
dedicated_llm._generate_summary = AsyncMock(return_value=("Dedicated summary", 5))
|
||||
|
||||
aggregator._summarizer.process_frame = capture_process
|
||||
context, config = self._create_context_and_config(dedicated_llm)
|
||||
original_message_count = len(context.messages)
|
||||
summarizer = LLMContextSummarizer(context=context, config=config)
|
||||
await summarizer.setup(self.task_manager)
|
||||
|
||||
# Call the method directly
|
||||
await aggregator._generate_summary_with_dedicated_llm(dedicated_llm, frame)
|
||||
# Track whether on_request_summarization event fires (it should NOT)
|
||||
event_fired = False
|
||||
|
||||
@summarizer.event_handler("on_request_summarization")
|
||||
async def on_request_summarization(summarizer, frame):
|
||||
nonlocal event_fired
|
||||
event_fired = True
|
||||
|
||||
# Trigger summarization via LLM response start
|
||||
from pipecat.frames.frames import LLMFullResponseStartFrame
|
||||
|
||||
await summarizer.process_frame(LLMFullResponseStartFrame())
|
||||
|
||||
# Wait for the background task to complete
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# The event should NOT have fired (dedicated LLM handles it internally)
|
||||
self.assertFalse(event_fired)
|
||||
|
||||
# Verify the dedicated LLM was called
|
||||
dedicated_llm._generate_summary.assert_called_once_with(frame)
|
||||
dedicated_llm._generate_summary.assert_called_once()
|
||||
|
||||
# Verify result was fed to the summarizer
|
||||
self.assertEqual(len(result_frames), 1)
|
||||
result = result_frames[0]
|
||||
self.assertIsInstance(result, LLMContextSummaryResultFrame)
|
||||
self.assertEqual(result.request_id, "dedicated_test")
|
||||
self.assertEqual(result.summary, "Dedicated summary")
|
||||
self.assertEqual(result.last_summarized_index, 1)
|
||||
self.assertIsNone(result.error)
|
||||
# Verify summary was applied to context (message count should decrease)
|
||||
self.assertLess(len(context.messages), original_message_count)
|
||||
|
||||
# Verify summary message is present
|
||||
summary_messages = [
|
||||
msg for msg in context.messages if "Conversation summary:" in msg.get("content", "")
|
||||
]
|
||||
self.assertEqual(len(summary_messages), 1)
|
||||
self.assertIn("Dedicated summary", summary_messages[0]["content"])
|
||||
|
||||
await summarizer.cleanup()
|
||||
|
||||
async def test_dedicated_llm_timeout(self):
|
||||
"""Test that dedicated LLM timeout produces error result."""
|
||||
from pipecat.processors.aggregators.llm_response_universal import (
|
||||
LLMAssistantAggregator,
|
||||
LLMAssistantAggregatorParams,
|
||||
)
|
||||
"""Test that dedicated LLM timeout produces error and clears state."""
|
||||
from pipecat.processors.aggregators.llm_context_summarizer import LLMContextSummarizer
|
||||
|
||||
context, _ = self._create_context_and_frame()
|
||||
|
||||
# Create a mock dedicated LLM that hangs
|
||||
dedicated_llm = LLMService()
|
||||
|
||||
async def slow_summary(frame):
|
||||
@@ -741,161 +735,116 @@ class TestDedicatedLLMSummarization(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
dedicated_llm._generate_summary = slow_summary
|
||||
|
||||
config = LLMContextSummarizationConfig(
|
||||
max_context_tokens=50,
|
||||
llm=dedicated_llm,
|
||||
)
|
||||
params = LLMAssistantAggregatorParams(
|
||||
enable_context_summarization=True,
|
||||
context_summarization_config=config,
|
||||
)
|
||||
aggregator = LLMAssistantAggregator(context, params=params)
|
||||
context, config = self._create_context_and_config(dedicated_llm)
|
||||
config.summarization_timeout = 0.1 # Very short timeout
|
||||
summarizer = LLMContextSummarizer(context=context, config=config)
|
||||
await summarizer.setup(self.task_manager)
|
||||
|
||||
# Mock summarizer.process_frame to capture the result
|
||||
result_frames = []
|
||||
original_message_count = len(context.messages)
|
||||
|
||||
async def capture_process(frame):
|
||||
result_frames.append(frame)
|
||||
# Trigger summarization
|
||||
from pipecat.frames.frames import LLMFullResponseStartFrame
|
||||
|
||||
aggregator._summarizer.process_frame = capture_process
|
||||
await summarizer.process_frame(LLMFullResponseStartFrame())
|
||||
|
||||
# Create frame with very short timeout
|
||||
frame = LLMContextSummaryRequestFrame(
|
||||
request_id="timeout_test",
|
||||
context=context,
|
||||
min_messages_to_keep=1,
|
||||
target_context_tokens=1000,
|
||||
summarization_prompt="Summarize this",
|
||||
summarization_timeout=0.1,
|
||||
)
|
||||
# Wait for the background task to complete (timeout + some buffer)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
await aggregator._generate_summary_with_dedicated_llm(dedicated_llm, frame)
|
||||
# Context should be unchanged (timeout = error = no summary applied)
|
||||
self.assertEqual(len(context.messages), original_message_count)
|
||||
|
||||
# Verify error result was fed to summarizer
|
||||
self.assertEqual(len(result_frames), 1)
|
||||
result = result_frames[0]
|
||||
self.assertIsInstance(result, LLMContextSummaryResultFrame)
|
||||
self.assertEqual(result.request_id, "timeout_test")
|
||||
self.assertEqual(result.summary, "")
|
||||
self.assertEqual(result.last_summarized_index, -1)
|
||||
self.assertIn("timed out", result.error)
|
||||
# Summarization state should be cleared so new requests can be made
|
||||
self.assertFalse(summarizer._summarization_in_progress)
|
||||
|
||||
await summarizer.cleanup()
|
||||
|
||||
async def test_dedicated_llm_exception(self):
|
||||
"""Test that dedicated LLM exceptions produce error result."""
|
||||
from pipecat.processors.aggregators.llm_response_universal import (
|
||||
LLMAssistantAggregator,
|
||||
LLMAssistantAggregatorParams,
|
||||
)
|
||||
"""Test that dedicated LLM exceptions produce error and clear state."""
|
||||
from pipecat.processors.aggregators.llm_context_summarizer import LLMContextSummarizer
|
||||
|
||||
context, frame = self._create_context_and_frame()
|
||||
|
||||
# Create a mock dedicated LLM that raises
|
||||
dedicated_llm = LLMService()
|
||||
dedicated_llm._generate_summary = AsyncMock(
|
||||
side_effect=RuntimeError("LLM connection failed")
|
||||
)
|
||||
|
||||
config = LLMContextSummarizationConfig(
|
||||
max_context_tokens=50,
|
||||
llm=dedicated_llm,
|
||||
)
|
||||
params = LLMAssistantAggregatorParams(
|
||||
enable_context_summarization=True,
|
||||
context_summarization_config=config,
|
||||
)
|
||||
aggregator = LLMAssistantAggregator(context, params=params)
|
||||
aggregator.push_error = AsyncMock()
|
||||
context, config = self._create_context_and_config(dedicated_llm)
|
||||
summarizer = LLMContextSummarizer(context=context, config=config)
|
||||
await summarizer.setup(self.task_manager)
|
||||
|
||||
# Mock summarizer.process_frame to capture the result
|
||||
result_frames = []
|
||||
original_message_count = len(context.messages)
|
||||
|
||||
async def capture_process(frame):
|
||||
result_frames.append(frame)
|
||||
# Trigger summarization
|
||||
from pipecat.frames.frames import LLMFullResponseStartFrame
|
||||
|
||||
aggregator._summarizer.process_frame = capture_process
|
||||
await summarizer.process_frame(LLMFullResponseStartFrame())
|
||||
|
||||
await aggregator._generate_summary_with_dedicated_llm(dedicated_llm, frame)
|
||||
# Wait for the background task to complete
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Verify error result was fed to summarizer
|
||||
self.assertEqual(len(result_frames), 1)
|
||||
result = result_frames[0]
|
||||
self.assertIsInstance(result, LLMContextSummaryResultFrame)
|
||||
self.assertEqual(result.request_id, "dedicated_test")
|
||||
self.assertEqual(result.summary, "")
|
||||
self.assertEqual(result.last_summarized_index, -1)
|
||||
self.assertIn("LLM connection failed", result.error)
|
||||
# Context should be unchanged (exception = error = no summary applied)
|
||||
self.assertEqual(len(context.messages), original_message_count)
|
||||
|
||||
# push_error should have been called
|
||||
aggregator.push_error.assert_called_once()
|
||||
# Summarization state should be cleared
|
||||
self.assertFalse(summarizer._summarization_in_progress)
|
||||
|
||||
async def test_on_request_summarization_routes_to_dedicated_llm(self):
|
||||
"""Test that _on_request_summarization routes to dedicated LLM when configured."""
|
||||
from pipecat.processors.aggregators.llm_response_universal import (
|
||||
LLMAssistantAggregator,
|
||||
LLMAssistantAggregatorParams,
|
||||
)
|
||||
await summarizer.cleanup()
|
||||
|
||||
context, frame = self._create_context_and_frame()
|
||||
async def test_dedicated_llm_does_not_emit_event(self):
|
||||
"""Test that summarizer does NOT emit on_request_summarization when dedicated LLM is set."""
|
||||
from pipecat.processors.aggregators.llm_context_summarizer import LLMContextSummarizer
|
||||
|
||||
dedicated_llm = LLMService()
|
||||
dedicated_llm._generate_summary = AsyncMock(return_value=("Summary", 1))
|
||||
|
||||
config = LLMContextSummarizationConfig(
|
||||
max_context_tokens=50,
|
||||
llm=dedicated_llm,
|
||||
)
|
||||
params = LLMAssistantAggregatorParams(
|
||||
enable_context_summarization=True,
|
||||
context_summarization_config=config,
|
||||
)
|
||||
aggregator = LLMAssistantAggregator(context, params=params)
|
||||
aggregator.push_frame = AsyncMock()
|
||||
context, config = self._create_context_and_config(dedicated_llm)
|
||||
summarizer = LLMContextSummarizer(context=context, config=config)
|
||||
await summarizer.setup(self.task_manager)
|
||||
|
||||
# Track what coroutine is passed to create_task
|
||||
created_coros = []
|
||||
original_create_task = aggregator.create_task
|
||||
event_fired = False
|
||||
|
||||
def mock_create_task(coro, *args, **kwargs):
|
||||
created_coros.append(coro)
|
||||
# Actually run the coroutine to avoid "never awaited" warning
|
||||
task = asyncio.ensure_future(coro)
|
||||
return task
|
||||
@summarizer.event_handler("on_request_summarization")
|
||||
async def on_request_summarization(summarizer, frame):
|
||||
nonlocal event_fired
|
||||
event_fired = True
|
||||
|
||||
aggregator.create_task = mock_create_task
|
||||
from pipecat.frames.frames import LLMFullResponseStartFrame
|
||||
|
||||
await aggregator._on_request_summarization(aggregator._summarizer, frame)
|
||||
await summarizer.process_frame(LLMFullResponseStartFrame())
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Should NOT push frame upstream
|
||||
aggregator.push_frame.assert_not_called()
|
||||
self.assertFalse(event_fired)
|
||||
|
||||
# Should have created a task for the dedicated LLM
|
||||
self.assertEqual(len(created_coros), 1)
|
||||
await summarizer.cleanup()
|
||||
|
||||
# Wait for the task to complete
|
||||
await asyncio.sleep(0.05)
|
||||
async def test_no_dedicated_llm_emits_event(self):
|
||||
"""Test that summarizer emits on_request_summarization when no dedicated LLM."""
|
||||
from pipecat.processors.aggregators.llm_context_summarizer import LLMContextSummarizer
|
||||
|
||||
async def test_on_request_summarization_pushes_upstream_without_dedicated_llm(self):
|
||||
"""Test that _on_request_summarization pushes upstream when no dedicated LLM."""
|
||||
from pipecat.processors.aggregators.llm_response_universal import (
|
||||
LLMAssistantAggregator,
|
||||
LLMAssistantAggregatorParams,
|
||||
)
|
||||
from pipecat.processors.frame_processor import FrameDirection
|
||||
|
||||
context, frame = self._create_context_and_frame()
|
||||
context = LLMContext()
|
||||
for i in range(10):
|
||||
context.add_message(
|
||||
{"role": "user", "content": f"Test message {i} that adds tokens to context."}
|
||||
)
|
||||
|
||||
config = LLMContextSummarizationConfig(max_context_tokens=50)
|
||||
params = LLMAssistantAggregatorParams(
|
||||
enable_context_summarization=True,
|
||||
context_summarization_config=config,
|
||||
)
|
||||
aggregator = LLMAssistantAggregator(context, params=params)
|
||||
aggregator.push_frame = AsyncMock()
|
||||
summarizer = LLMContextSummarizer(context=context, config=config)
|
||||
await summarizer.setup(self.task_manager)
|
||||
|
||||
await aggregator._on_request_summarization(aggregator._summarizer, frame)
|
||||
request_frame = None
|
||||
|
||||
# Should push frame upstream
|
||||
aggregator.push_frame.assert_called_once_with(frame, FrameDirection.UPSTREAM)
|
||||
@summarizer.event_handler("on_request_summarization")
|
||||
async def on_request_summarization(summarizer, frame):
|
||||
nonlocal request_frame
|
||||
request_frame = frame
|
||||
|
||||
from pipecat.frames.frames import LLMFullResponseStartFrame
|
||||
|
||||
await summarizer.process_frame(LLMFullResponseStartFrame())
|
||||
|
||||
self.assertIsNotNone(request_frame)
|
||||
self.assertIsInstance(request_frame, LLMContextSummaryRequestFrame)
|
||||
|
||||
await summarizer.cleanup()
|
||||
|
||||
|
||||
class TestLLMSpecificMessageHandling(unittest.TestCase):
|
||||
|
||||
Reference in New Issue
Block a user