Update app.py

app.py

@@ -670,8 +670,12 @@ class BraveSearchClient:
                 return []
 
 
+# Global state for web search settings
+web_search_settings = {}
+
+
 class OpenAIHandler(AsyncStreamHandler):
-    def __init__(self, web_search_enabled: bool = False, search_client: Optional[BraveSearchClient] = None, webrtc_id: str = None) -> None:
+    def __init__(self) -> None:
         super().__init__(
             expected_layout="mono",
             output_sample_rate=SAMPLE_RATE,
@@ -680,16 +684,15 @@ class OpenAIHandler(AsyncStreamHandler):
         )
         self.connection = None
         self.output_queue = asyncio.Queue()
-        self.web_search_enabled = web_search_enabled
-        self.search_client = search_client
+        self.search_client = None
         self.function_call_in_progress = False
         self.current_function_args = ""
         self.current_call_id = None
-        self.webrtc_id = webrtc_id
+        self.webrtc_id = None
+        self.web_search_enabled = False
 
     def copy(self):
-        # Return a new instance with the same settings
-        return OpenAIHandler(self.web_search_enabled, self.search_client, self.webrtc_id)
+        return OpenAIHandler()
 
     async def search_web(self, query: str) -> str:
         """Perform web search and return formatted results"""
@@ -714,6 +717,10 @@ class OpenAIHandler(AsyncStreamHandler):
 
     async def start_up(self):
         """Connect to realtime API with function calling enabled"""
+        # Get web search setting from global state
+        if self.webrtc_id and self.webrtc_id in web_search_settings:
+            self.web_search_enabled = web_search_settings[self.webrtc_id]
+        
         print(f"Starting up handler with web_search_enabled={self.web_search_enabled}")
         self.client = openai.AsyncOpenAI()
         
@@ -724,7 +731,7 @@ class OpenAIHandler(AsyncStreamHandler):
                 "type": "function",
                 "function": {
                     "name": "web_search",
-                    "description": "Search the web for information",
+                    "description": "Search the web for information when the user asks questions about current events, news, or topics that require up-to-date information",
                     "parameters": {
                         "type": "object",
                         "properties": {
@@ -749,6 +756,16 @@ class OpenAIHandler(AsyncStreamHandler):
                 "tool_choice": "auto" if tools else "none"
             }
             
+            # Add instructions to use web search
+            if self.web_search_enabled:
+                session_update["instructions"] = (
+                    "You are a helpful assistant. When users ask about current events, "
+                    "news, recent information, or topics that might have changed recently, "
+                    "use the web_search function to find the most up-to-date information. "
+                    "Always search when the user asks about: weather, news, current events, "
+                    "recent updates, prices, or any time-sensitive information."
+                )
+            
             await conn.session.update(session=session_update)
             self.connection = conn
             print(f"Connected with tools: {len(tools)} functions")
@@ -769,7 +786,7 @@ class OpenAIHandler(AsyncStreamHandler):
                 
                 # Handle function calls
                 elif event.type == "response.function_call_arguments.start":
-                    print(f"Function call started, call_id: {getattr(event, 'call_id', 'unknown')}")
+                    print(f"Function call started")
                     self.function_call_in_progress = True
                     self.current_function_args = ""
                     self.current_call_id = getattr(event, 'call_id', None)
@@ -830,9 +847,6 @@ class OpenAIHandler(AsyncStreamHandler):
             self.connection = None
 
 
-# Store handlers by webrtc_id
-handlers = {}
-
 # Initialize search client
 brave_api_key = os.getenv("BSEARCH_API")
 search_client = BraveSearchClient(brave_api_key) if brave_api_key else None
@@ -842,65 +856,73 @@ def update_chatbot(chatbot: list[dict], response: ResponseAudioTranscriptDoneEve
     chatbot.append({"role": "assistant", "content": response.transcript})
     return chatbot
 
+# Create chatbot component
+chatbot = gr.Chatbot(type="messages")
+
+# Create base handler and stream
+handler = OpenAIHandler()
+handler.search_client = search_client
+
+stream = Stream(
+    handler,
+    mode="send-receive",
+    modality="audio",
+    additional_inputs=[chatbot],
+    additional_outputs=[chatbot],
+    additional_outputs_handler=update_chatbot,
+    rtc_configuration=get_twilio_turn_credentials() if get_space() else None,
+    concurrency_limit=5 if get_space() else None,
+    time_limit=90 if get_space() else None,
+)
 
-# Create base stream without handler for mounting
 app = FastAPI()
+stream.mount(app)
 
-@app.post("/webrtc/offer")
-async def handle_offer(request: dict):
-    """Handle WebRTC offer and create a new handler for this connection"""
-    webrtc_id = request.get("webrtc_id")
-    web_search_enabled = request.get("web_search_enabled", False)
-    
-    print(f"Received offer - webrtc_id: {webrtc_id}, web_search_enabled: {web_search_enabled}")
-    
-    # Create a new handler for this connection
-    handler = OpenAIHandler(
-        web_search_enabled=web_search_enabled,
-        search_client=search_client,
-        webrtc_id=webrtc_id
-    )
-    
-    # Store the handler
-    handlers[webrtc_id] = handler
-    
-    # Create chatbot instance for this connection
-    chatbot = gr.Chatbot(type="messages")
-    
-    # Create stream for this connection
-    stream = Stream(
-        handler,
-        mode="send-receive",
-        modality="audio",
-        additional_inputs=[chatbot],
-        additional_outputs=[chatbot],
-        additional_outputs_handler=update_chatbot,
-        rtc_configuration=get_twilio_turn_credentials() if get_space() else None,
-        concurrency_limit=5 if get_space() else None,
-        time_limit=90 if get_space() else None,
-    )
+# Override the offer handler to capture web search settings
+original_offer_handler = stream.offer
+
+async def custom_offer_handler(request):
+    """Custom offer handler to process web search settings"""
+    # Extract web search setting from request
+    if hasattr(request, 'web_search_enabled'):
+        web_search_enabled = request.web_search_enabled
+    elif isinstance(request, dict) and 'web_search_enabled' in request:
+        web_search_enabled = request['web_search_enabled']
+    else:
+        web_search_enabled = False
     
-    # Mount stream temporarily
-    stream.mount(app)
+    # Extract webrtc_id
+    if hasattr(request, 'webrtc_id'):
+        webrtc_id = request.webrtc_id
+    elif isinstance(request, dict) and 'webrtc_id' in request:
+        webrtc_id = request['webrtc_id']
+    else:
+        webrtc_id = None
     
-    # Store stream reference
-    handler.stream = stream
+    print(f"Received offer - webrtc_id: {webrtc_id}, web_search_enabled: {web_search_enabled}")
     
-    # Process the offer
-    response = await stream.offer(request)
+    # Store web search setting
+    if webrtc_id:
+        web_search_settings[webrtc_id] = web_search_enabled
+        handler.webrtc_id = webrtc_id
+        handler.web_search_enabled = web_search_enabled
+        
+        # Restart handler if already connected to update settings
+        if handler.connection:
+            await handler.shutdown()
+            await handler.start_up()
     
-    return response
+    # Call original offer handler
+    return await original_offer_handler(request)
 
+# Replace the offer method
+stream.offer = custom_offer_handler
 
 @app.get("/outputs")
 async def outputs(webrtc_id: str):
     """Stream outputs including search events"""
     async def output_stream():
-        handler = handlers.get(webrtc_id)
-        if not handler or not hasattr(handler, 'stream'):
-            return
-        
-        async for output in handler.stream.output_stream(webrtc_id):
+        async for output in stream.output_stream(webrtc_id):
             if hasattr(output, 'args') and output.args:
                 # Check if it's a search event
                 if isinstance(output.args[0], dict) and output.args[0].get('type') == 'search':
@@ -926,38 +948,8 @@ if __name__ == "__main__":
     
     mode = os.getenv("MODE")
     if mode == "UI":
-        # For UI mode, we need a base handler and stream
-        base_handler = OpenAIHandler()
-        chatbot = gr.Chatbot(type="messages")
-        stream = Stream(
-            base_handler,
-            mode="send-receive",
-            modality="audio",
-            additional_inputs=[chatbot],
-            additional_outputs=[chatbot],
-            additional_outputs_handler=update_chatbot,
-            rtc_configuration=get_twilio_turn_credentials() if get_space() else None,
-            concurrency_limit=5 if get_space() else None,
-            time_limit=90 if get_space() else None,
-        )
-        stream.mount(app)
         stream.ui.launch(server_port=7860)
     elif mode == "PHONE":
-        # Similar for phone mode
-        base_handler = OpenAIHandler()
-        chatbot = gr.Chatbot(type="messages")
-        stream = Stream(
-            base_handler,
-            mode="send-receive",
-            modality="audio",
-            additional_inputs=[chatbot],
-            additional_outputs=[chatbot],
-            additional_outputs_handler=update_chatbot,
-            rtc_configuration=get_twilio_turn_credentials() if get_space() else None,
-            concurrency_limit=5 if get_space() else None,
-            time_limit=90 if get_space() else None,
-        )
-        stream.mount(app)
         stream.fastphone(host="0.0.0.0", port=7860)
     else:
         uvicorn.run(app, host="0.0.0.0", port=7860)