Spaces:
Running
Running
| # handler.py | |
| import asyncio | |
| import base64 | |
| import json | |
| import os | |
| import traceback | |
| from websockets.asyncio.client import connect | |
| host = "generativelanguage.googleapis.com" | |
| model = "gemini-2.0-flash-exp" | |
| api_key = os.environ["GOOGLE_API_KEY"] | |
| uri = f"wss://{host}/ws/google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContent?key={api_key}" | |
| class AudioLoop: | |
| def __init__(self): | |
| self.ws = None | |
| # Queue for messages to be sent *to* Gemini | |
| self.out_queue = asyncio.Queue() | |
| # Queue for PCM audio received *from* Gemini | |
| self.audio_in_queue = asyncio.Queue() | |
| async def startup(self, tools=None): | |
| """Send the model setup message to Gemini. | |
| Args: | |
| tools: Optional list of tools to enable for the model | |
| """ | |
| setup_msg = {"setup": {"model": f"models/{model}"}} | |
| if tools: | |
| setup_msg["setup"]["tools"] = tools | |
| await self.ws.send(json.dumps(setup_msg)) | |
| raw_response = await self.ws.recv() | |
| setup_response = json.loads(raw_response) | |
| print("[AudioLoop] Setup response from Gemini:", setup_response) | |
| async def send_realtime(self): | |
| """Read from out_queue and forward those messages to Gemini in real time.""" | |
| while True: | |
| msg = await self.out_queue.get() | |
| await self.ws.send(json.dumps(msg)) | |
| async def receive_audio(self): | |
| """Read from Gemini websocket and push PCM data into audio_in_queue.""" | |
| async for raw_response in self.ws: | |
| response = json.loads(raw_response) | |
| # Debug log all responses (optional) | |
| # print("Gemini raw response:", response) | |
| # Check if there's inline PCM data | |
| try: | |
| b64data = ( | |
| response["serverContent"]["modelTurn"]["parts"][0]["inlineData"]["data"] | |
| ) | |
| pcm_data = base64.b64decode(b64data) | |
| await self.audio_in_queue.put(pcm_data) | |
| except KeyError: | |
| # No audio in this message | |
| pass | |
| tool_call = response.pop('toolCall', None) | |
| if tool_call is not None: | |
| await self.handle_tool_call(tool_call) | |
| # If "turnComplete" is present | |
| if "serverContent" in response and response["serverContent"].get("turnComplete"): | |
| print("[AudioLoop] Gemini turn complete") | |
| async def handle_tool_call(self,tool_call): | |
| print(" ", tool_call) | |
| for fc in tool_call['functionCalls']: | |
| msg = { | |
| 'tool_response': { | |
| 'function_responses': [{ | |
| 'id': fc['id'], | |
| 'name': fc['name'], | |
| 'response':{'result': {'string_value': 'ok'}} | |
| }] | |
| } | |
| } | |
| print('>>> ', msg) | |
| await self.ws.send(json.dumps(msg)) | |
| async def run(self): | |
| """Main entry point: connects to Gemini, starts send/receive tasks.""" | |
| try: | |
| turn_on_the_lights_schema = {'name': 'turn_on_the_lights'} | |
| turn_off_the_lights_schema = {'name': 'turn_off_the_lights'} | |
| tools = [ | |
| {'google_search': {}}, | |
| {'function_declarations': [turn_on_the_lights_schema, turn_off_the_lights_schema]}, | |
| {'code_execution': {}}, | |
| ] | |
| async with connect(uri, additional_headers={"Content-Type": "application/json"}) as ws: | |
| self.ws = ws | |
| await self.startup(tools) | |
| try: | |
| async with asyncio.TaskGroup() as tg: | |
| send_task = tg.create_task(self.send_realtime()) | |
| receive_task = tg.create_task(self.receive_audio()) | |
| await asyncio.Future() # Keep running until canceled | |
| finally: | |
| # Clean up tasks and connection | |
| if 'send_task' in locals(): | |
| send_task.cancel() | |
| if 'receive_task' in locals(): | |
| receive_task.cancel() | |
| try: | |
| await self.ws.close() | |
| print("[AudioLoop] Closed WebSocket connection") | |
| except Exception as e: | |
| print(f"[AudioLoop] Error closing Gemini connection: {e}") | |
| print("[AudioLoop] Cleanup complete") | |
| except asyncio.CancelledError: | |
| print("[AudioLoop] Cancelled") | |
| except Exception as e: | |
| traceback.print_exc() | |
| raise | |