Update webapp.py

webapp.py

@@ -1,30 +1,31 @@
-# updated_webapp.py
+# basic_webapp.py
 
 import asyncio
 import base64
 import json
 import os
-from typing import Optional, Dict, Any, List
+import sys
 
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 import uvicorn
 
-from handler import AudioLoop  # Import AudioLoop class
+# Import the simplified AudioLoop
+from basic_handler import AudioLoop
 
 app = FastAPI()
 
 # Store active client connections
-client_connections: Dict[str, Dict[str, Any]] = {}
+active_clients = {}
 
-# Mount the web_ui directory to serve static files
+# Mount static files directory
 current_dir = os.path.dirname(os.path.realpath(__file__))
 app.mount("/static", StaticFiles(directory=current_dir), name="static")
 
 @app.get("/")
 async def get_index():
-    # Read and return the index.html file
+    """Serve the main HTML interface."""
     index_path = os.path.join(current_dir, "index.html")
     with open(index_path, "r", encoding="utf-8") as f:
         html_content = f.read()
@@ -32,6 +33,7 @@ async def get_index():
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
+    """Handle WebSocket connections from clients."""
     await websocket.accept()
     print("[websocket_endpoint] Client connected.")
     
@@ -40,47 +42,45 @@ async def websocket_endpoint(websocket: WebSocket):
     
     # Create a new AudioLoop instance for this client
     audio_loop = AudioLoop()
-    client_connections[client_id] = {
+    active_clients[client_id] = {
         "websocket": websocket,
         "audio_loop": audio_loop,
-        "audio_ordering_buffer": {},
-        "expected_audio_seq": 0,
-        "repo_url": None,
-        "preferences": None
+        "repo_context": None
     }
     
     # Start the AudioLoop for this client
     loop_task = asyncio.create_task(audio_loop.run())
-    print(f"[websocket_endpoint] Started new AudioLoop for client {client_id}")
+    print(f"[websocket_endpoint] Started AudioLoop for client {client_id}")
 
-    async def from_client_to_gemini():
-        """Handles incoming messages from the client and forwards them to Gemini."""
+    async def process_client_messages():
+        """Handle messages from the client and forward to Gemini."""
         try:
             while True:
                 data = await websocket.receive_text()
                 msg = json.loads(data)
-                msg_type = msg.get("type")
+                msg_type = msg.get("type", "")
 
-                # Handle repository URL and preferences
                 if msg_type == "init":
-                    client_connections[client_id]["repo_url"] = msg.get("repo_url", "")
-                    client_connections[client_id]["preferences"] = {
+                    # Store repository context info
+                    active_clients[client_id]["repo_context"] = {
+                        "repo_url": msg.get("repo_url", ""),
                         "github_token": msg.get("github_token", ""),
                         "user_type": msg.get("user_type", "coder"),
                         "response_detail": msg.get("response_detail", "normal")
                     }
-                    print(f"[from_client_to_gemini] Client {client_id} initialized with repo: {client_connections[client_id]['repo_url']}")
-                    # Send a confirmation back to client
+                    print(f"[process_client_messages] Stored context for {client_id}: {msg.get('repo_url', '')}")
+                    
+                    # Send confirmation
                     await websocket.send_text(json.dumps({
                         "type": "status",
                         "status": "initialized",
-                        "message": "G.E.N.I.E. is ready to assist with this repository."
+                        "message": "Ready to assist with this repository."
                     }))
-
-                # Handle audio data from client
+                
                 elif msg_type == "audio":
-                    # Decode base64 audio data
+                    # Forward audio data to Gemini
                     raw_pcm = base64.b64decode(msg["payload"])
+                    
                     forward_msg = {
                         "realtime_input": {
                             "media_chunks": [
@@ -91,48 +91,27 @@ async def websocket_endpoint(websocket: WebSocket):
                             ]
                         }
                     }
-                    
-                    # Retrieve the sequence number from the message
-                    seq = msg.get("seq")
-                    audio_ordering_buffer = client_connections[client_id]["audio_ordering_buffer"]
-                    expected_audio_seq = client_connections[client_id]["expected_audio_seq"]
-                    
-                    if seq is not None:
-                        # Store the message in the buffer
-                        audio_ordering_buffer[seq] = forward_msg
-                        # Forward any messages in order
-                        while expected_audio_seq in audio_ordering_buffer:
-                            msg_to_forward = audio_ordering_buffer.pop(expected_audio_seq)
-                            await audio_loop.out_queue.put(msg_to_forward)
-                            expected_audio_seq += 1
-                        client_connections[client_id]["expected_audio_seq"] = expected_audio_seq
-                    else:
-                        # If no sequence number is provided, forward immediately
-                        await audio_loop.out_queue.put(forward_msg)
-
-                # Handle text data from client
+                    await audio_loop.out_queue.put(forward_msg)
+                
                 elif msg_type == "text":
+                    # Process text query from client
                     user_text = msg.get("content", "")
-                    # Augment the query with repository context if available
-                    repo_context = ""
-                    if client_connections[client_id]["repo_url"]:
-                        repo_context = f"For GitHub repository: {client_connections[client_id]['repo_url']}"
-                        
-                        # Add preferences context
-                        prefs = client_connections[client_id]["preferences"]
-                        if prefs:
-                            user_type = prefs.get("user_type", "coder")
-                            detail_level = prefs.get("response_detail", "normal")
-                            repo_context += f"\nUser role: {user_type}, Preferred detail level: {detail_level}"
                     
-                    # Combine context with user query
-                    if repo_context:
-                        enhanced_text = f"{repo_context}\n\nUser query: {user_text}"
-                    else:
-                        enhanced_text = user_text
+                    # Add repository context if available
+                    context = active_clients[client_id]["repo_context"]
+                    if context and context["repo_url"]:
+                        # Format context info for Gemini
+                        context_text = (
+                            f"The GitHub repository being discussed is: {context['repo_url']}\n"
+                            f"User role: {context['user_type']}\n"
+                            f"Preferred detail level: {context['response_detail']}\n\n"
+                            f"Please consider this context when answering the following question:\n"
+                        )
+                        user_text = context_text + user_text
                     
-                    print(f"[from_client_to_gemini] Forwarding user text to Gemini: {enhanced_text[:100]}...")
+                    print(f"[process_client_messages] Sending text to Gemini: {user_text[:100]}...")
                     
+                    # Format message for Gemini
                     forward_msg = {
                         "client_content": {
                             "turn_complete": True,
@@ -140,81 +119,98 @@ async def websocket_endpoint(websocket: WebSocket):
                                 {
                                     "role": "user",
                                     "parts": [
-                                        {"text": enhanced_text}
+                                        {"text": user_text}
                                     ]
                                 }
                             ]
                         }
                     }
                     await audio_loop.out_queue.put(forward_msg)
-
-                # Handle interrupt request
+                
                 elif msg_type == "interrupt":
-                    print(f"[from_client_to_gemini] Client {client_id} requested interrupt")
-                    # TODO: Send interrupt signal to Gemini if possible
-                    # For now, just acknowledge the interrupt request
+                    # For now, just acknowledge the interrupt
+                    # This is a simple implementation because true interruption
+                    # may require additional API support
+                    print(f"[process_client_messages] Interrupt requested by {client_id}")
                     await websocket.send_text(json.dumps({
                         "type": "status",
                         "status": "interrupted",
-                        "message": "G.E.N.I.E. processing interrupted by user."
+                        "message": "Processing interrupted by user."
                     }))
-
+                
                 else:
-                    print(f"[from_client_to_gemini] Unknown message type: {msg_type}")
+                    print(f"[process_client_messages] Unknown message type: {msg_type}")
 
         except WebSocketDisconnect:
-            print(f"[from_client_to_gemini] Client {client_id} disconnected.")
+            print(f"[process_client_messages] Client {client_id} disconnected")
             cleanup_client(client_id, loop_task)
         except Exception as e:
-            print(f"[from_client_to_gemini] Error: {e}")
+            print(f"[process_client_messages] Error: {e}")
             cleanup_client(client_id, loop_task)
 
-    async def from_gemini_to_client():
-        """Reads PCM audio from Gemini and sends it back to the client."""
+    async def forward_gemini_responses():
+        """Read responses from Gemini and send them to the client."""
         try:
             while True:
-                pcm_data = await audio_loop.audio_in_queue.get()
-                b64_pcm = base64.b64encode(pcm_data).decode()
-
-                out_msg = {
-                    "type": "audio",
-                    "payload": b64_pcm
-                }
-                print(f"[from_gemini_to_client] Sending audio chunk to client {client_id}. Size: {len(pcm_data)}")
-                await websocket.send_text(json.dumps(out_msg))
+                # Check for audio data
+                try:
+                    pcm_data = await asyncio.wait_for(audio_loop.audio_in_queue.get(), 0.5)
+                    b64_pcm = base64.b64encode(pcm_data).decode()
+                    
+                    # Send audio to client
+                    out_msg = {
+                        "type": "audio",
+                        "payload": b64_pcm
+                    }
+                    print(f"[forward_gemini_responses] Sending audio chunk to client {client_id}")
+                    await websocket.send_text(json.dumps(out_msg))
+                except asyncio.TimeoutError:
+                    # No audio available, continue checking
+                    pass
+                
+                # We could add additional processing for text responses here
+                # if we had a separate queue for text content
 
         except WebSocketDisconnect:
-            print(f"[from_gemini_to_client] Client {client_id} disconnected.")
+            print(f"[forward_gemini_responses] Client {client_id} disconnected")
             cleanup_client(client_id, loop_task)
         except Exception as e:
-            print(f"[from_gemini_to_client] Error: {e}")
+            print(f"[forward_gemini_responses] Error: {e}")
             cleanup_client(client_id, loop_task)
 
-    def cleanup_client(client_id, loop_task):
-        """Clean up resources for a disconnected client."""
-        if client_id in client_connections:
-            # Cancel the AudioLoop task
-            if loop_task and not loop_task.done():
-                loop_task.cancel()
-            # Remove the client from active connections
-            del client_connections[client_id]
-            print(f"[cleanup_client] Cleaned up resources for client {client_id}")
-
-    # Launch both tasks concurrently. If either fails or disconnects, we exit.
+    def cleanup_client(client_id, task):
+        """Clean up resources when a client disconnects."""
+        if client_id in active_clients:
+            client_data = active_clients[client_id]
+            
+            # Stop the AudioLoop
+            if "audio_loop" in client_data:
+                client_data["audio_loop"].stop()
+            
+            # Cancel the task if it's still running
+            if task and not task.done():
+                task.cancel()
+            
+            # Remove from active clients
+            del active_clients[client_id]
+            print(f"[cleanup_client] Cleaned up resources for {client_id}")
+
+    # Run both tasks concurrently
     try:
         await asyncio.gather(
-            from_client_to_gemini(),
-            from_gemini_to_client(),
+            process_client_messages(),
+            forward_gemini_responses()
         )
     finally:
-        print(f"[websocket_endpoint] WebSocket handler finished for client {client_id}.")
+        print(f"[websocket_endpoint] WebSocket handler finished for {client_id}")
         cleanup_client(client_id, loop_task)
 
 if __name__ == "__main__":
-    # Make sure the GOOGLE_API_KEY environment variable is set before running
+    # Verify API key is present
     if "GOOGLE_API_KEY" not in os.environ:
         print("Error: GOOGLE_API_KEY environment variable not set")
-        print("Please set it with: export GOOGLE_API_KEY='your_api_key'")
-        exit(1)
-        
+        print("Please set it with: export GOOGLE_API_KEY='your_api_key_here'")
+        sys.exit(1)
+    
+    # Start the server
     uvicorn.run("webapp:app", host="0.0.0.0", port=7860, reload=True)