Spaces:
Sleeping
Sleeping
File size: 6,591 Bytes
f6f98ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
import asyncio
import base64
import json
import os
import traceback
from typing import Dict, List, Any, Optional, Tuple
from handler import AudioLoop
import gradio as gr
class GeminiGradioIntegration:
"""
Integration class that connects the Gemini AudioLoop with Gradio components
"""
def __init__(self):
self.audio_loop = None
self.response_queue = asyncio.Queue()
self.is_connected = False
async def initialize(self):
"""Initialize the connection to Gemini"""
try:
if self.audio_loop is None:
self.audio_loop = AudioLoop()
# Start the AudioLoop
self.loop_task = asyncio.create_task(self.audio_loop.run())
# Start response listener
self.response_task = asyncio.create_task(self.listen_for_responses())
self.is_connected = True
return "Connected to Gemini successfully"
else:
return "Already connected to Gemini"
except Exception as e:
self.is_connected = False
traceback.print_exc()
return f"Failed to connect to Gemini: {str(e)}"
async def listen_for_responses(self):
"""Listen for responses from Gemini's audio queue and process them"""
try:
while True:
# Get PCM data from the audio_loop
pcm_data = await self.audio_loop.audio_in_queue.get()
# Put the data in our response queue for Gradio to pick up
await self.response_queue.put({
"type": "audio",
"data": pcm_data
})
# In a real implementation, we'd also listen for text responses
# and tool calls from Gemini, process them, and update the UI
except asyncio.CancelledError:
print("[GeminiGradioIntegration] Response listener cancelled")
except Exception as e:
print(f"[GeminiGradioIntegration] Error in response listener: {str(e)}")
traceback.print_exc()
async def send_query(self, repo_url: str, query: str) -> str:
"""
Send a text query to Gemini about a GitHub repository
"""
if not self.is_connected:
return "Not connected to Gemini. Please connect first."
try:
# Format the query to include the repository context
formatted_query = f"Analyze GitHub repository: {repo_url}\nQuery: {query}"
# Send the query to Gemini
forward_msg = {
"client_content": {
"turn_complete": True,
"turns": [
{
"role": "user",
"parts": [
{"text": formatted_query}
]
}
]
}
}
await self.audio_loop.out_queue.put(forward_msg)
return "Query sent to Gemini"
except Exception as e:
traceback.print_exc()
return f"Error sending query: {str(e)}"
async def send_audio(self, audio_data):
"""
Send audio data to Gemini
"""
if not self.is_connected:
return "Not connected to Gemini. Please connect first."
try:
# Convert the audio data to PCM format
# Note: In a real implementation, you'd need to convert from Gradio's
# audio format to the format expected by Gemini
# For simplicity, let's assume audio_data is already in PCM format
raw_pcm = audio_data.tobytes()
# Create the message to send to Gemini
forward_msg = {
"realtime_input": {
"media_chunks": [
{
"data": base64.b64encode(raw_pcm).decode(),
"mime_type": "audio/pcm"
}
]
}
}
# Send to Gemini
await self.audio_loop.out_queue.put(forward_msg)
return "Audio sent to Gemini"
except Exception as e:
traceback.print_exc()
return f"Error sending audio: {str(e)}"
async def get_next_response(self):
"""
Get the next response from Gemini (for Gradio to poll)
"""
try:
# Non-blocking check if there's a response
if self.response_queue.empty():
return None
# Get the response
response = await self.response_queue.get()
return response
except Exception as e:
print(f"[GeminiGradioIntegration] Error getting response: {str(e)}")
return None
def cleanup(self):
"""Clean up resources when shutting down"""
if self.audio_loop:
# Cancel the tasks
if hasattr(self, 'loop_task') and self.loop_task:
self.loop_task.cancel()
if hasattr(self, 'response_task') and self.response_task:
self.response_task.cancel()
self.is_connected = False
print("[GeminiGradioIntegration] Cleaned up resources")
# Initialize the integration
gemini_integration = GeminiGradioIntegration()
# Function to be called from Gradio
def connect_to_gemini():
"""Connect to Gemini (for Gradio button)"""
return asyncio.run(gemini_integration.initialize())
def send_text_query(repo_url, query):
"""Send a text query to Gemini (for Gradio interface)"""
return asyncio.run(gemini_integration.send_query(repo_url, query))
def send_audio_query(audio_data):
"""Send audio to Gemini (for Gradio interface)"""
if audio_data is None:
return "No audio data provided"
return asyncio.run(gemini_integration.send_audio(audio_data))
# This function would be called periodically by Gradio to check for new responses
def check_for_responses():
"""Check for new responses from Gemini"""
response = asyncio.run(gemini_integration.get_next_response())
if response and response["type"] == "audio":
# Convert PCM data back to a format Gradio can understand
# This depends on how Gradio expects audio data
return response["data"]
return None
|