LLMAssistantContextAggregator: create a task to run on_context_updated

This commit is contained in:
Aleix Conchillo Flaqué
2025-03-24 15:39:35 -07:00
parent 3b3fdd0da1
commit 397bae29f7
3 changed files with 19 additions and 5 deletions

View File

@@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Fixed
- Fixed an issue that would cause `LLMAssistantContextAggregator` to block
processing more frames while processing a function call result.
- Fixed an issue where the `RTVIObserver` would report two bot started and
stopped speaking events for each bot turn.

View File

@@ -6,7 +6,7 @@
import asyncio
from abc import abstractmethod
from typing import Dict, List
from typing import Dict, List, Set
from loguru import logger
@@ -380,6 +380,7 @@ class LLMAssistantContextAggregator(LLMContextResponseAggregator):
self._started = 0
self._function_calls_in_progress: Dict[str, FunctionCallInProgressFrame] = {}
self._context_updated_tasks: Set[asyncio.Task] = ()
async def handle_aggregation(self, aggregation: str):
self._context.add_message({"role": "assistant", "content": aggregation})
@@ -486,10 +487,13 @@ class LLMAssistantContextAggregator(LLMContextResponseAggregator):
if run_llm:
await self.push_context_frame(FrameDirection.UPSTREAM)
# Emit the on_context_updated callback once the function call
# result is added to the context
# Call the `on_context_updated` callback once the function call result
# is added to the context. Also, run this in a separate task to make
# sure we don't block the pipeline.
if properties and properties.on_context_updated:
await properties.on_context_updated()
task = await self.create_task(properties.on_context_updated())
self._context_updated_tasks.add(task)
task.add_done_callback(self._context_updated_task_finished)
async def _handle_function_call_cancel(self, frame: FunctionCallCancelFrame):
logger.debug(
@@ -535,6 +539,13 @@ class LLMAssistantContextAggregator(LLMContextResponseAggregator):
else:
self._aggregation += frame.text
def _context_updated_task_finished(self, task: asyncio.Task):
self._context_updated_tasks.discard(task)
# The task is finished so this should exit immediately. We need to do
# this because otherwise the task manager would report a dangling task
# if we don't remove it.
asyncio.run_coroutine_threadsafe(self.wait_for_task(task), self.get_event_loop())
class LLMUserResponseAggregator(LLMUserContextAggregator):
def __init__(self, messages: List[dict] = [], **kwargs):

View File

@@ -369,7 +369,7 @@ class LLMService(AIService):
if tuple_to_remove:
self._function_call_tasks.discard(tuple_to_remove)
# The task is finished so this should exit immediately. We need to
# do this because otherwise the task manager would have a dangling
# do this because otherwise the task manager would report a dangling
# task if we don't remove it.
asyncio.run_coroutine_threadsafe(self.wait_for_task(task), self.get_event_loop())