loop - restart connection - gemini

[pgrayy]

Description

Restart Gemini model connection after timeout.

Note, Gemini uses a live session handle id to restart connections, which means we do not need to seed the message history. Regardless, we do not have a message history recorded yet for Gemini because of googleapis/python-genai#1504. With that said, right now the Gemini connection cannot be Strands session persisted. We will address this in a follow up.

For additional context, the session handle id can be used to restart connections over a 24 hour period. This means if we were to go longer, we would at that point need to switch from the handle id to using a message history. In either case, we need to work on session managing both the live handle id and the message history.

Testing

• I ran hatch run bidi:prepare: Wrote new unit tests
• I ran the following test script:

import asyncio
import json

from strands import tool
from strands.experimental.bidi import BidiAgent
from strands.experimental.bidi.io import BidiAudioIO, BidiTextIO
from strands.experimental.bidi.models import BidiGeminiLiveModel
from strands.experimental.bidi.types import BidiConnectionRestartEvent, BidiOutputEvent
from strands.experimental.hooks.events import BidiBeforeConnectionRestartEvent
from strands.experimental.hooks.events import BidiAfterConnectionRestartEvent
from strands.hooks import HookProvider, HookRegistry


class RestartHook(HookProvider):
    def register_hooks(self, registry: HookRegistry) -> None:
        registry.add_callback(BidiBeforeConnectionRestartEvent, self.print_before)
        registry.add_callback(BidiAfterConnectionRestartEvent, self.print_after)

    def print_before(self, event: BidiBeforeConnectionRestartEvent) -> None:
        print(f"event=<{event}> | before hook called")

    def print_after(self, event: BidiAfterConnectionRestartEvent) -> None:
        print(f"event=<{event}> | after hook called")


async def print_restart(event: BidiOutputEvent) -> None:
    if isinstance(event, BidiConnectionRestartEvent):
        print(f"event=<{event}> | restart typed event")


@tool
async def time_tool() -> str:
    print("tool=<time>")
    return "12:01"


async def main() -> None:
    print("MAIN - starting agent")
    model = BidiGeminiLiveModel(...)
    agent = BidiAgent(model=model, hooks=[RestartHook()], tools=[time_tool])

    audio_io = BidiAudioIO()
    text_io = BidiTextIO()

    try:
        run_coro = agent.run(inputs=[audio_io.input()], outputs=[audio_io.output(), text_io.output(), print_restart])
        await asyncio.wait_for(run_coro, timeout=60)
    except asyncio.TimeoutError:
        pass

    print(f"MAIN - stopping agent: {json.dumps(agent.messages, indent=2)}")


if __name__ == "__main__":
    asyncio.run(main())

Gemini successfully restarted. Additionally, Gemini remembered the conversation from the previous connection (restarts every 15 minutes). Gemini however was not able to remember the tool use requested in the previous connection. It has this weird behavior where it will continue to request the tool until it receives the result. We will be sure to document this odd behavior. Nothing much we can do on our end.

[pgrayy]

Removing as we will want to be explicit with what arguments we unpack from kwargs and place into config.