apunkt/dograh
1
0
Fork 0
mirror of https://github.com/dograh-hq/dograh.git synced 2026-07-25 12:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity

fix: fix transition logic for realtime providers

Browse source
This commit is contained in:
Abhishek Kumar 2026-07-15 16:45:16 +05:30
parent 348cd8427b
commit 960d183c82
7 changed files with 415 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

192
api/services/pipecat/realtime/gemini_live.py
View file

@ -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

6
api/services/workflow/pipecat_engine.py
View file

@ -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]

9
api/services/workflow/pipecat_engine_custom_tools.py
View file

@ -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(

36
api/tests/test_custom_tools.py
View file

@ -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."""

202
api/tests/test_gemini_live_reconnect_tool_results.py
View file

@ -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()

1
api/tests/test_pipecat_engine_tool_calls.py
View file

@ -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(

2
pipecat

@ -1 +1 @@
Subproject commit f4feab0c59129702b3a6e6f818f62229cc663f1a
Subproject commit fed8541fca8893b318b862779f1962ad9903de81