From 960d183c820df25b7edbc3f1d7c91582edfab564 Mon Sep 17 00:00:00 2001 From: Abhishek Kumar Date: Wed, 15 Jul 2026 16:45:16 +0530 Subject: [PATCH] fix: fix transition logic for realtime providers --- api/services/pipecat/realtime/gemini_live.py | 192 ++++++++++++++--- api/services/workflow/pipecat_engine.py | 6 +- .../workflow/pipecat_engine_custom_tools.py | 9 + api/tests/test_custom_tools.py | 36 ++++ ...test_gemini_live_reconnect_tool_results.py | 202 +++++++++++++++++- api/tests/test_pipecat_engine_tool_calls.py | 1 + pipecat | 2 +- 7 files changed, 415 insertions(+), 33 deletions(-) diff --git a/api/services/pipecat/realtime/gemini_live.py b/api/services/pipecat/realtime/gemini_live.py index 31218b85..6b85a4a4 100644 --- a/api/services/pipecat/realtime/gemini_live.py +++ b/api/services/pipecat/realtime/gemini_live.py @@ -9,8 +9,8 @@ Layers Dograh engine integration quirks onto upstream-pristine - **Reconnect on node transitions.** Gemini Live cannot update ``system_instruction`` mid-session, so a setting change triggers a reconnect (deferred until the bot turn ends if currently responding). -- **Function-call deferral.** Tool calls emitted mid-turn are queued and run - when the bot stops speaking, to avoid racing the turn's audio. +- **Node-transition deferral.** Node-transition calls emitted mid-turn are + queued and run when the bot stops speaking, to avoid cutting off its audio. - **User-mute audio gating.** ``UserMuteStarted/StoppedFrame`` from the user aggregator gates whether incoming audio is forwarded to Gemini. - **TTSSpeakFrame as greeting trigger.** The engine queues a TTSSpeakFrame @@ -18,6 +18,7 @@ Layers Dograh engine integration quirks onto upstream-pristine it and runs the initial-context path. """ +import asyncio from typing import Any from google.genai.types import Content, Part @@ -44,6 +45,11 @@ from pipecat.utils.tracing.service_decorators import traced_gemini_live class DograhGeminiLiveLLMService(GeminiLiveLLMService): """Gemini Live with Dograh engine integration quirks. See module docstring.""" + # Gemini input transcription is delivered independently from tool calls. + # Give late transcription messages a small window to arrive before running + # a node-transition function and tearing down the current Live connection. + _NODE_TRANSITION_TRANSCRIPTION_GRACE_SECONDS = 0.5 + # Route tool schemas through Gemini's ``parameters_json_schema`` field so # MCP/imported tools that use JSON Schema keywords (``const``, ``not``, # nested ``anyOf``) rejected by the strict ``Schema`` model are accepted. @@ -59,15 +65,19 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService): # Guards initial-response triggering against double-firing across the # initial TTSSpeakFrame and any LLMContextFrame that may arrive. self._handled_initial_context: bool = False - # When a system_instruction change arrives mid-bot-turn, the reconnect - # is queued and drained when the turn ends. - self._reconnect_pending: bool = False - # Function calls emitted by Gemini mid-bot-turn are deferred here and - # invoked when the turn ends, so they don't race the turn's audio. - self._pending_function_calls: list[FunctionCallFromLLM] = [] + # Node-transition calls emitted mid-bot-turn are deferred here so the + # transition does not tear down Gemini while it is still producing audio. + self._pending_node_transition_function_calls: list[FunctionCallFromLLM] = [] # Text greeting captured from the first TTSSpeakFrame while the Gemini # session is still connecting. self._pending_initial_greeting_text: str | None = None + self._transition_function_call_task: asyncio.Task | None = None + # Intentional node changes use a fresh, context-seeded connection rather + # than a potentially stale session-resumption handle. The new connection + # remains gated until the function-call result has landed in LLMContext. + self._awaiting_node_transition_context: bool = False + self._node_transition_context_received: bool = False + self._node_transition_context_seed_started: bool = False # ------------------------------------------------------------------ # Hooks from upstream GeminiLiveLLMService @@ -78,33 +88,109 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService): # lets pre-call fetch populate template variables first. return bool(self._settings.system_instruction) + def _requires_node_transition_context_aggregation(self) -> bool: + # A node transition replaces the current Gemini Live connection and + # seeds the new one from our local LLMContext. Wait for the upstream + # user aggregator to commit any final TranscriptionFrame before + # set_node() changes the prompt and starts that reconnect. + return True + + async def cleanup(self) -> None: + """Cancel a delayed transition before tearing down the Live session.""" + if self._transition_function_call_task: + await self.cancel_task(self._transition_function_call_task) + self._transition_function_call_task = None + await super().cleanup() + async def _handle_changed_settings(self, changed: dict[str, Any]) -> set[str]: if "system_instruction" not in changed: return set() + + # PipecatEngine updates system_instruction only from set_node(). The + # first set_node happens before a Live session exists; every later one + # is a node transition whose tool call has already been deferred until + # the current bot turn finishes. if not self._session: # First-time setting after deferred-connect. await self._connect() - elif self._bot_is_responding: - # Bot is mid-turn — drain the reconnect when it ends so we don't - # cut the bot off mid-utterance. - self._reconnect_pending = True else: - await self._reconnect() + await self._reconnect_for_node_transition() return {"system_instruction"} async def _run_or_defer_function_calls( self, function_calls_llm: list[FunctionCallFromLLM] ): + if not self._contains_node_transition(function_calls_llm): + await super()._run_or_defer_function_calls(function_calls_llm) + return + + # Keep a provider tool-call batch together. Splitting a mixed batch here + # would discard Pipecat's shared function-call group and could trigger an + # LLM run before every result from the original batch has arrived. if self._bot_is_responding: # Latest batch wins; Gemini emits tool calls as one batch per # tool_call message, so this overwrite is intentional. - self._pending_function_calls = function_calls_llm + self._pending_node_transition_function_calls = function_calls_llm logger.debug( - f"{self}: deferring {len(function_calls_llm)} function call(s) " + f"{self}: deferring {len(function_calls_llm)} node-transition " + "function call(s) " "until bot turn ends" ) return - await super()._run_or_defer_function_calls(function_calls_llm) + + self._schedule_node_transition_function_calls(function_calls_llm) + + def _contains_node_transition( + self, function_calls_llm: list[FunctionCallFromLLM] + ) -> bool: + return any(self._is_node_transition(fc) for fc in function_calls_llm) + + def _is_node_transition(self, function_call: FunctionCallFromLLM) -> bool: + return self._function_is_node_transition(function_call.function_name) + + def _schedule_node_transition_function_calls( + self, function_calls_llm: list[FunctionCallFromLLM] + ) -> None: + """Run transition calls after late input transcription has settled.""" + if ( + self._transition_function_call_task + and not self._transition_function_call_task.done() + ): + logger.warning( + f"{self}: node-transition function call already pending; " + "ignoring duplicate batch" + ) + return + + async def _run_after_transcription_grace() -> None: + try: + await asyncio.sleep(self._NODE_TRANSITION_TRANSCRIPTION_GRACE_SECONDS) + await self._flush_pending_user_transcription() + await self.run_function_calls(function_calls_llm) + finally: + self._transition_function_call_task = None + + self._transition_function_call_task = self.create_task( + _run_after_transcription_grace(), + name=f"{self}::node-transition-function-calls", + ) + + async def _flush_pending_user_transcription(self) -> None: + """Publish any punctuationless user transcript before a node handoff.""" + if self._transcription_timeout_task: + if not self._transcription_timeout_task.done(): + await self.cancel_task(self._transcription_timeout_task) + self._transcription_timeout_task = None + + if not self._user_transcription_buffer: + return + + text = self._user_transcription_buffer + self._user_transcription_buffer = "" + logger.debug( + f"{self}: flushing pending user transcription before node transition" + ) + await self._push_user_transcription(text, result=None) # ------------------------------------------------------------------ # State-transition side effects @@ -114,22 +200,35 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService): was_responding = self._bot_is_responding await super()._set_bot_is_responding(responding) if was_responding and not responding: - await self._run_pending_function_calls() - if self._reconnect_pending: - self._reconnect_pending = False - await self._reconnect() + await self._run_pending_node_transition_function_calls() - async def _run_pending_function_calls(self): - """Run any function calls deferred during the bot's last turn.""" - if not self._pending_function_calls: + async def _run_pending_node_transition_function_calls(self): + """Run any node-transition calls deferred during the bot's last turn.""" + if not self._pending_node_transition_function_calls: return - fcs = self._pending_function_calls - self._pending_function_calls = [] + fcs = self._pending_node_transition_function_calls + self._pending_node_transition_function_calls = [] logger.debug( - f"{self}: executing {len(fcs)} deferred function call(s) " + f"{self}: executing {len(fcs)} deferred node-transition call(s) " "after bot turn ended" ) - await self.run_function_calls(fcs) + self._schedule_node_transition_function_calls(fcs) + + async def _reconnect_for_node_transition(self) -> None: + """Start a fresh connection and wait to seed the completed context. + + Gemini can report ``resumable=False`` while generating or executing a + function call. A workflow transition happens at exactly that boundary, + so using the last (older) resumption handle can omit the triggering user + turn. Use the local LLMContext as the source of truth for this intentional + handoff instead. + """ + self._awaiting_node_transition_context = True + self._node_transition_context_received = False + self._node_transition_context_seed_started = False + self._session_resumption_handle = None + await self._disconnect() + await self._connect(session_resumption_handle=None) # ------------------------------------------------------------------ # Frame handling: mute, TTSSpeakFrame, BotStoppedSpeakingFrame flush @@ -165,9 +264,9 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService): if isinstance(frame, BotStoppedSpeakingFrame): # Belt-and-suspenders: the main drain happens in # _set_bot_is_responding(False), but if Gemini delays turn_complete - # past the audible end of the turn, flushing here ensures pending - # function calls fire promptly. - await self._run_pending_function_calls() + # past the audible end of the turn, flushing here ensures a pending + # node transition fires promptly. + await self._run_pending_node_transition_function_calls() # Fall through to super for the actual push. await super().process_frame(frame, direction) @@ -185,6 +284,11 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService): # ------------------------------------------------------------------ async def _handle_context(self, context: LLMContext): + if self._awaiting_node_transition_context: + self._context = context + self._node_transition_context_received = True + await self._maybe_seed_node_transition_context() + return if not self._handled_initial_context: self._handled_initial_context = True self._context = context @@ -249,6 +353,12 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService): f"In _handle_session_ready self._run_llm_when_session_ready: {self._run_llm_when_session_ready}" ) self._session = session + if self._awaiting_node_transition_context: + # Do not accept realtime input until the function-call result frame + # has updated the shared context and that complete history is seeded. + self._ready_for_realtime_input = False + await self._maybe_seed_node_transition_context() + return self._ready_for_realtime_input = True if self._run_llm_when_session_ready: # Context arrived before session was ready — fulfil the queued @@ -266,3 +376,25 @@ class DograhGeminiLiveLLMService(GeminiLiveLLMService): # a handle (e.g. node transitions before any handle was issued) are # followed by a function-call-result LLMContextFrame which feeds the # updated-context branch in _handle_context. + + async def _maybe_seed_node_transition_context(self) -> None: + if ( + not self._awaiting_node_transition_context + or not self._node_transition_context_received + or not self._session + or self._node_transition_context_seed_started + ): + return + + self._node_transition_context_seed_started = True + try: + # The complete tool result is already present in the history being + # seeded, so mark it delivered locally instead of sending a provider + # tool response for a call that the fresh session never issued. + await self._process_completed_function_calls(send_new_results=False) + await self._create_initial_response() + self._awaiting_node_transition_context = False + self._node_transition_context_received = False + await self._drain_pending_tool_results() + finally: + self._node_transition_context_seed_started = False diff --git a/api/services/workflow/pipecat_engine.py b/api/services/workflow/pipecat_engine.py index 716a2800..423e0e41 100644 --- a/api/services/workflow/pipecat_engine.py +++ b/api/services/workflow/pipecat_engine.py @@ -344,7 +344,11 @@ class PipecatEngine: ) # Register function with LLM - self.llm.register_function(name, transition_func) + self.llm.register_function( + name, + transition_func, + is_node_transition=True, + ) async def _register_knowledge_base_function( self, document_uuids: list[str] diff --git a/api/services/workflow/pipecat_engine_custom_tools.py b/api/services/workflow/pipecat_engine_custom_tools.py index 2a39568d..6e544e06 100644 --- a/api/services/workflow/pipecat_engine_custom_tools.py +++ b/api/services/workflow/pipecat_engine_custom_tools.py @@ -288,10 +288,19 @@ class CustomToolManager: # Create and register the handler handler, timeout_secs = self._create_handler(tool, function_name) + # End-call and transfer-call tools are workflow-control + # boundaries even though they do not necessarily select another + # graph node. Give them the same ordering guarantees as an + # explicit node-transition function. + is_node_transition = tool.category in { + ToolCategory.END_CALL.value, + ToolCategory.TRANSFER_CALL.value, + } self._engine.llm.register_function( function_name, handler, timeout_secs=timeout_secs, + is_node_transition=is_node_transition, ) logger.debug( diff --git a/api/tests/test_custom_tools.py b/api/tests/test_custom_tools.py index bfa7652b..5ff818c3 100644 --- a/api/tests/test_custom_tools.py +++ b/api/tests/test_custom_tools.py @@ -1415,6 +1415,7 @@ class TestCustomToolManagerUnit: # Verify handler was registered assert "api_call" in registered_handlers assert registered_kwargs["api_call"]["timeout_secs"] == pytest.approx(5) + assert registered_kwargs["api_call"]["is_node_transition"] is False # Now test that the handler works handler = registered_handlers["api_call"] @@ -1445,6 +1446,41 @@ class TestCustomToolManagerUnit: # Verify result was returned assert result_received["status"] == "success" + @pytest.mark.asyncio + @pytest.mark.parametrize("category", ["end_call", "transfer_call"]) + async def test_register_handlers_marks_call_control_tools_as_node_transitions( + self, category + ): + from api.services.workflow.pipecat_engine_custom_tools import CustomToolManager + + mock_engine = Mock() + mock_engine._get_organization_id = AsyncMock(return_value=1) + mock_engine.llm.register_function = Mock() + manager = CustomToolManager(mock_engine) + tool = MockToolModel( + tool_uuid=f"{category}-uuid", + name=category.replace("_", " "), + description=f"Perform {category}", + category=category, + definition={ + "schema_version": 1, + "type": category, + "config": {}, + }, + ) + + with patch( + "api.services.workflow.pipecat_engine_custom_tools.db_client.get_tools_by_uuids", + new=AsyncMock(return_value=[tool]), + ): + await manager.register_handlers([tool.tool_uuid]) + + mock_engine.llm.register_function.assert_called_once() + assert ( + mock_engine.llm.register_function.call_args.kwargs["is_node_transition"] + is True + ) + @pytest.mark.asyncio async def test_transfer_call_renders_destination_from_initial_context(self): """Transfer call tools resolve destination templates before provider calls.""" diff --git a/api/tests/test_gemini_live_reconnect_tool_results.py b/api/tests/test_gemini_live_reconnect_tool_results.py index 75ca8f35..42031d2b 100644 --- a/api/tests/test_gemini_live_reconnect_tool_results.py +++ b/api/tests/test_gemini_live_reconnect_tool_results.py @@ -1,11 +1,20 @@ +import asyncio import json from types import SimpleNamespace from unittest.mock import AsyncMock import pytest -from pipecat.frames.frames import TranscriptionFrame, TTSSpeakFrame +from pipecat.frames.frames import ( + NodeTransitionStartedFrame, + TranscriptionFrame, + TTSSpeakFrame, +) from pipecat.processors.aggregators.llm_context import LLMContext +from pipecat.processors.aggregators.llm_response_universal import ( + LLMContextAggregatorPair, +) from pipecat.processors.frame_processor import FrameDirection +from pipecat.services.llm_service import FunctionCallFromLLM from api.services.pipecat.realtime.gemini_live import DograhGeminiLiveLLMService @@ -163,3 +172,194 @@ async def test_tts_greeting_waits_for_gemini_session_before_sending_prompt(): assert prompt.endswith('"Hello from Dograh."') assert service._run_llm_when_session_ready is False assert service._pending_initial_greeting_text is None + + +@pytest.mark.asyncio +async def test_transition_call_flushes_pending_transcription_before_execution(): + service = _make_service() + service._NODE_TRANSITION_TRANSCRIPTION_GRACE_SECONDS = 0 + service.register_function( + "transition_to_next_node", + AsyncMock(), + is_node_transition=True, + ) + service._user_transcription_buffer = "My last answer" + + events = [] + + async def _push_transcription(text, result=None): + events.append(("transcription", text)) + + async def _run_function_calls(function_calls): + events.append(("function", function_calls[0].function_name)) + + service._push_user_transcription = AsyncMock(side_effect=_push_transcription) + service.run_function_calls = AsyncMock(side_effect=_run_function_calls) + service.create_task = lambda coro, name=None: asyncio.create_task(coro, name=name) + + function_call = FunctionCallFromLLM( + context=LLMContext(), + tool_call_id="call-transition", + function_name="transition_to_next_node", + arguments={}, + ) + + await service._run_or_defer_function_calls([function_call]) + transition_task = service._transition_function_call_task + assert transition_task is not None + await transition_task + + assert events == [ + ("transcription", "My last answer"), + ("function", "transition_to_next_node"), + ] + assert service._user_transcription_buffer == "" + + +@pytest.mark.asyncio +async def test_non_transition_call_is_not_deferred_while_bot_is_responding(): + service = _make_service() + service._bot_is_responding = True + service.run_function_calls = AsyncMock() + + function_call = FunctionCallFromLLM( + context=LLMContext(), + tool_call_id="call-ordinary", + function_name="look_up_account", + arguments={}, + ) + + await service._run_or_defer_function_calls([function_call]) + + service.run_function_calls.assert_awaited_once_with([function_call]) + assert service._pending_node_transition_function_calls == [] + assert service._transition_function_call_task is None + + +@pytest.mark.asyncio +async def test_node_transition_call_is_deferred_while_bot_is_responding(): + service = _make_service() + service._bot_is_responding = True + service.register_function( + "transition_to_next_node", + AsyncMock(), + is_node_transition=True, + ) + service.run_function_calls = AsyncMock() + + function_call = FunctionCallFromLLM( + context=LLMContext(), + tool_call_id="call-transition", + function_name="transition_to_next_node", + arguments={}, + ) + + await service._run_or_defer_function_calls([function_call]) + + service.run_function_calls.assert_not_awaited() + assert service._pending_node_transition_function_calls == [function_call] + assert service._transition_function_call_task is None + + +@pytest.mark.asyncio +async def test_initial_system_instruction_update_is_not_a_node_handoff(): + service = _make_service() + service._connect = AsyncMock() + service._disconnect = AsyncMock() + + handled = await service._handle_changed_settings( + {"system_instruction": "old prompt"} + ) + + assert handled == {"system_instruction"} + service._connect.assert_awaited_once_with() + service._disconnect.assert_not_awaited() + assert service._awaiting_node_transition_context is False + + +@pytest.mark.asyncio +async def test_node_transition_uses_fresh_connection_instead_of_stale_handle(): + service = _make_service() + service._session = _FakeSession() + service._session_resumption_handle = "stale-handle" + service._disconnect = AsyncMock() + service._connect = AsyncMock() + + handled = await service._handle_changed_settings( + {"system_instruction": "old prompt"} + ) + + assert handled == {"system_instruction"} + assert service._session_resumption_handle is None + assert service._awaiting_node_transition_context is True + service._disconnect.assert_awaited_once() + service._connect.assert_awaited_once_with(session_resumption_handle=None) + + +@pytest.mark.asyncio +async def test_fresh_transition_session_waits_for_updated_context_before_ready(): + service = _make_service() + service._handled_initial_context = True + service._awaiting_node_transition_context = True + service._node_transition_context_received = False + service._process_completed_function_calls = AsyncMock() + service._drain_pending_tool_results = AsyncMock() + + async def _seed_context(): + service._ready_for_realtime_input = True + + service._create_initial_response = AsyncMock(side_effect=_seed_context) + + session = _FakeSession() + await service._handle_session_ready(session) + + assert service._ready_for_realtime_input is False + service._create_initial_response.assert_not_awaited() + + context = _make_tool_result_context("call-transition") + await service._handle_context(context) + + service._process_completed_function_calls.assert_awaited_once_with( + send_new_results=False + ) + service._create_initial_response.assert_awaited_once() + service._drain_pending_tool_results.assert_awaited_once() + assert service._ready_for_realtime_input is True + assert service._awaiting_node_transition_context is False + + +@pytest.mark.asyncio +async def test_node_transition_frame_commits_user_transcript_to_context(): + context = LLMContext() + context_aggregator = LLMContextAggregatorPair(context, realtime_service_mode=True) + user_aggregator = context_aggregator.user() + user_aggregator.push_context_frame = AsyncMock() + + await user_aggregator._handle_transcription( + TranscriptionFrame( + text="The answer that selected this node", + user_id="", + timestamp="now", + ) + ) + context_aggregation_event = asyncio.Event() + await user_aggregator._handle_node_transition_started( + NodeTransitionStartedFrame( + function_calls=[ + FunctionCallFromLLM( + context=context, + tool_call_id="call-transition", + function_name="transition_to_next_node", + arguments={}, + ) + ], + context_aggregation_event=context_aggregation_event, + ) + ) + + assert context.messages[-1] == { + "role": "user", + "content": "The answer that selected this node", + } + user_aggregator.push_context_frame.assert_awaited_once() + assert context_aggregation_event.is_set() diff --git a/api/tests/test_pipecat_engine_tool_calls.py b/api/tests/test_pipecat_engine_tool_calls.py index 92d3c54d..3a92659d 100644 --- a/api/tests/test_pipecat_engine_tool_calls.py +++ b/api/tests/test_pipecat_engine_tool_calls.py @@ -182,6 +182,7 @@ class TestPipecatEngineToolCalls: # Assert that the context was updated with END_CALL_SYSTEM_PROMPT assert llm._settings.system_instruction == END_CALL_SYSTEM_PROMPT + assert llm._functions["end_call"].is_node_transition is True @pytest.mark.asyncio async def test_parallel_builtin_and_transition_calls_through_engine_1( diff --git a/pipecat b/pipecat index f4feab0c..fed8541f 160000 --- a/pipecat +++ b/pipecat @@ -1 +1 @@ -Subproject commit f4feab0c59129702b3a6e6f818f62229cc663f1a +Subproject commit fed8541fca8893b318b862779f1962ad9903de81