Update app.py

app.py

@@ -670,14 +670,17 @@ class BraveSearchClient:
                 return []
 
 
-# Global state for web search settings
+# Initialize search client globally
+brave_api_key = os.getenv("BSEARCH_API")
+search_client = BraveSearchClient(brave_api_key) if brave_api_key else None
+print(f"Search client initialized: {search_client is not None}, API key present: {bool(brave_api_key)}")
+
+# Store web search settings by connection
 web_search_settings = {}
-current_webrtc_id = None
-current_web_search_enabled = False
 
 
 class OpenAIHandler(AsyncStreamHandler):
-    def __init__(self) -> None:
+    def __init__(self, web_search_enabled: bool = False, webrtc_id: str = None) -> None:
         super().__init__(
             expected_layout="mono",
             output_sample_rate=SAMPLE_RATE,
@@ -686,21 +689,23 @@ class OpenAIHandler(AsyncStreamHandler):
         )
         self.connection = None
         self.output_queue = asyncio.Queue()
-        self.search_client = search_client  # Use global search client
+        self.search_client = search_client
         self.function_call_in_progress = False
         self.current_function_args = ""
         self.current_call_id = None
-        self.webrtc_id = None
-        self.web_search_enabled = False
-        self.startup_called = False
+        self.webrtc_id = webrtc_id
+        self.web_search_enabled = web_search_enabled
+        print(f"Handler created with web_search_enabled={web_search_enabled}, webrtc_id={webrtc_id}")
 
     def copy(self):
-        # Create a new instance that inherits settings
-        new_handler = OpenAIHandler()
-        new_handler.search_client = self.search_client
-        new_handler.webrtc_id = current_webrtc_id  # Use global variable
-        new_handler.web_search_enabled = current_web_search_enabled  # Use global variable
-        return new_handler
+        # Check for stored settings when copying
+        if self.webrtc_id and self.webrtc_id in web_search_settings:
+            web_search_enabled = web_search_settings[self.webrtc_id]
+        else:
+            web_search_enabled = self.web_search_enabled
+        
+        print(f"Handler.copy() called - creating new handler with web_search_enabled={web_search_enabled}")
+        return OpenAIHandler(web_search_enabled=web_search_enabled, webrtc_id=self.webrtc_id)
 
     async def search_web(self, query: str) -> str:
         """Perform web search and return formatted results"""
@@ -725,28 +730,25 @@ class OpenAIHandler(AsyncStreamHandler):
 
     async def start_up(self):
         """Connect to realtime API with function calling enabled"""
-        # Check web search setting before connecting
-        if self.webrtc_id and self.webrtc_id in web_search_settings:
-            self.web_search_enabled = web_search_settings[self.webrtc_id]
-            print(f"Retrieved web_search_enabled={self.web_search_enabled} for webrtc_id={self.webrtc_id}")
-        
         print(f"Starting up handler with web_search_enabled={self.web_search_enabled}")
         self.client = openai.AsyncOpenAI()
         
         # Define the web search function
         tools = []
+        instructions = "You are a helpful assistant. Respond in Korean when the user speaks Korean."
+        
         if self.web_search_enabled and self.search_client:
             tools = [{
                 "type": "function",
                 "function": {
                     "name": "web_search",
-                    "description": "Search the web for information when the user asks questions about current events, news, weather, prices, or topics that require up-to-date information",
+                    "description": "Search the web for current information. Use this for weather, news, prices, current events, or any time-sensitive topics.",
                     "parameters": {
                         "type": "object",
                         "properties": {
                             "query": {
                                 "type": "string",
-                                "description": "The search query"
+                                "description": "The search query in Korean or English"
                             }
                         },
                         "required": ["query"]
@@ -754,6 +756,20 @@ class OpenAIHandler(AsyncStreamHandler):
                 }
             }]
             print("Web search function added to tools")
+            
+            instructions = (
+                "You are a helpful assistant with web search capabilities. "
+                "IMPORTANT: You MUST use the web_search function for ANY of these topics:\n"
+                "- Weather (날씨, 기온, 비, 눈)\n"
+                "- News (뉴스, 소식)\n"
+                "- Current events (현재, 최근, 오늘, 지금)\n"
+                "- Prices (가격, 환율, 주가)\n"
+                "- Sports scores or results\n"
+                "- Any question about 2024 or 2025\n"
+                "- Any time-sensitive information\n\n"
+                "When in doubt, USE web_search. It's better to search and provide accurate information "
+                "than to guess or use outdated information. Always respond in Korean when the user speaks Korean."
+            )
         
         async with self.client.beta.realtime.connect(
             model="gpt-4o-mini-realtime-preview-2024-12-17"
@@ -761,26 +777,20 @@ class OpenAIHandler(AsyncStreamHandler):
             # Update session with tools
             session_update = {
                 "turn_detection": {"type": "server_vad"},
+                "instructions": instructions,
                 "tools": tools,
                 "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, weather, 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), "
-                    "최근(recent), 최신(latest), 오늘(today), or any time-sensitive information."
-                )
-            
             await conn.session.update(session=session_update)
             self.connection = conn
             print(f"Connected with tools: {len(tools)} functions")
-            self.startup_called = True
             
             async for event in self.connection:
+                # Debug logging for function calls
+                if event.type.startswith("response.function_call"):
+                    print(f"Function event: {event.type}")
+                
                 if event.type == "response.audio_transcript.done":
                     await self.output_queue.put(AdditionalOutputs(event))
                 
@@ -843,10 +853,14 @@ class OpenAIHandler(AsyncStreamHandler):
     async def receive(self, frame: tuple[int, np.ndarray]) -> None:
         if not self.connection:
             return
-        _, array = frame
-        array = array.squeeze()
-        audio_message = base64.b64encode(array.tobytes()).decode("utf-8")
-        await self.connection.input_audio_buffer.append(audio=audio_message)
+        try:
+            _, array = frame
+            array = array.squeeze()
+            audio_message = base64.b64encode(array.tobytes()).decode("utf-8")
+            await self.connection.input_audio_buffer.append(audio=audio_message)
+        except Exception as e:
+            print(f"Error in receive: {e}")
+            # Connection might be closed, ignore the error
 
     async def emit(self) -> tuple[int, np.ndarray] | AdditionalOutputs | None:
         return await wait_for_item(self.output_queue)
@@ -857,24 +871,23 @@ class OpenAIHandler(AsyncStreamHandler):
             self.connection = None
 
 
-# Initialize search client
-brave_api_key = os.getenv("BSEARCH_API")
-search_client = BraveSearchClient(brave_api_key) if brave_api_key else None
-print(f"Search client initialized: {search_client is not None}, API key present: {bool(brave_api_key)}")
+# Factory function to create handler with settings
+def create_handler():
+    """Factory function to create handler - will be customized per connection"""
+    return OpenAIHandler(web_search_enabled=False)
+
 
 def update_chatbot(chatbot: list[dict], response: ResponseAudioTranscriptDoneEvent):
     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
+# Create components
+chatbot = gr.Chatbot(type="messages")
 
+# Create initial stream with factory function
 stream = Stream(
-    handler,
+    create_handler,  # Pass factory function instead of instance
     mode="send-receive",
     modality="audio",
     additional_inputs=[chatbot],
@@ -887,45 +900,60 @@ stream = Stream(
 
 app = FastAPI()
 
-# Mount the stream
-stream.mount(app)
+# Custom handler factory that reads settings
+def custom_handler_factory():
+    """Custom factory that checks for web search settings"""
+    # Try to get the most recent settings
+    if web_search_settings:
+        # Get the most recent webrtc_id
+        recent_id = max(web_search_settings.keys(), key=lambda k: web_search_settings[k].get('timestamp', 0))
+        settings = web_search_settings[recent_id]
+        return OpenAIHandler(
+            web_search_enabled=settings.get('enabled', False),
+            webrtc_id=recent_id
+        )
+    return OpenAIHandler(web_search_enabled=False)
 
-# Intercept the POST body before it reaches the stream
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.requests import Request as StarletteRequest
-import io
+# Replace the handler factory
+stream.handler = custom_handler_factory
 
-class WebSearchMiddleware(BaseHTTPMiddleware):
-    async def dispatch(self, request: StarletteRequest, call_next):
-        if request.url.path == "/webrtc/offer" and request.method == "POST":
-            # Read the body
-            body = await request.body()
-            
-            # Parse it
-            try:
-                data = json.loads(body)
-                webrtc_id = data.get("webrtc_id")
-                web_search_enabled = data.get("web_search_enabled", False)
-                
-                print(f"Middleware intercepted - webrtc_id: {webrtc_id}, web_search_enabled: {web_search_enabled}")
-                
-                # Store settings globally
-                global current_webrtc_id, current_web_search_enabled
-                if webrtc_id:
-                    current_webrtc_id = webrtc_id
-                    current_web_search_enabled = web_search_enabled
-                    web_search_settings[webrtc_id] = web_search_enabled
-            except:
-                pass
-            
-            # Recreate the request with the body
-            request._body = body
-            
-        response = await call_next(request)
-        return response
+# Mount stream
+stream.mount(app)
+
+# Intercept offer to capture settings
+@app.post("/webrtc/offer", include_in_schema=False)
+async def custom_offer(request: Request):
+    """Intercept offer to capture web search settings"""
+    body = await request.json()
+    
+    webrtc_id = body.get("webrtc_id")
+    web_search_enabled = body.get("web_search_enabled", False)
+    
+    print(f"Custom offer - webrtc_id: {webrtc_id}, web_search_enabled: {web_search_enabled}")
+    
+    # Store settings with timestamp
+    if webrtc_id:
+        web_search_settings[webrtc_id] = {
+            'enabled': web_search_enabled,
+            'timestamp': asyncio.get_event_loop().time()
+        }
+    
+    # Remove our custom route temporarily
+    custom_route = None
+    for i, route in enumerate(app.routes):
+        if hasattr(route, 'path') and route.path == "/webrtc/offer" and route.endpoint == custom_offer:
+            custom_route = app.routes.pop(i)
+            break
+    
+    # Forward to stream's offer handler
+    response = await stream.offer(body)
+    
+    # Re-add our custom route
+    if custom_route:
+        app.routes.insert(0, custom_route)
+    
+    return response
 
-# Add middleware
-app.add_middleware(WebSearchMiddleware)
 
 @app.get("/outputs")
 async def outputs(webrtc_id: str):