Private-AI

Running

App Files Files Community

seawolf2357 commited on 24 days ago

Commit

e4a5e76

verified ·

1 Parent(s): a3f3e10

Update app.py

Browse files

Files changed (1) hide show

app.py +98 -74

app.py CHANGED Viewed

@@ -397,18 +397,7 @@ HTML_CONTENT = """<!DOCTYPE html>
         searchToggle.addEventListener('click', () => {
             webSearchEnabled = !webSearchEnabled;
             searchToggle.classList.toggle('active', webSearchEnabled);
-            // Update server-side settings if connected
-            if (webrtc_id) {
-                fetch('/update-search-setting', {
-                    method: 'POST',
-                    headers: { 'Content-Type': 'application/json' },
-                    body: JSON.stringify({
-                        webrtc_id: webrtc_id,
-                        web_search_enabled: webSearchEnabled
-                    })
-                });
-            }
         });
         function updateStatus(state) {
@@ -577,7 +566,7 @@ HTML_CONTENT = """<!DOCTYPE html>
                 });
                 eventSource.addEventListener("search", (event) => {
                     const eventJson = JSON.parse(event.data);
-                    if (eventJson.results) {
                         addMessage("search-result", `웹 검색 중: "${eventJson.query}"`);
                     }
                 });
@@ -682,7 +671,7 @@ class BraveSearchClient:
 class OpenAIHandler(AsyncStreamHandler):
-    def __init__(self) -> None:
         super().__init__(
             expected_layout="mono",
             output_sample_rate=SAMPLE_RATE,
@@ -691,20 +680,23 @@ class OpenAIHandler(AsyncStreamHandler):
         )
         self.connection = None
         self.output_queue = asyncio.Queue()
-        self.web_search_enabled = False
-        self.search_client = None
         self.function_call_in_progress = False
         self.current_function_args = ""
-        self.call_id = None
     def copy(self):
-        return OpenAIHandler()
     async def search_web(self, query: str) -> str:
         """Perform web search and return formatted results"""
         if not self.search_client or not self.web_search_enabled:
             return "웹 검색이 비활성화되어 있습니다."
         results = await self.search_client.search(query)
         if not results:
             return f"'{query}'에 대한 검색 결과를 찾을 수 없습니다."
@@ -722,6 +714,7 @@ class OpenAIHandler(AsyncStreamHandler):
     async def start_up(self):
         """Connect to realtime API with function calling enabled"""
         self.client = openai.AsyncOpenAI()
         # Define the web search function
@@ -744,6 +737,7 @@ class OpenAIHandler(AsyncStreamHandler):
                     }
                 }
             }]
         async with self.client.beta.realtime.connect(
             model="gpt-4o-mini-realtime-preview-2024-12-17"
@@ -757,6 +751,7 @@ class OpenAIHandler(AsyncStreamHandler):
             await conn.session.update(session=session_update)
             self.connection = conn
             async for event in self.connection:
                 if event.type == "response.audio_transcript.done":
@@ -774,9 +769,10 @@ class OpenAIHandler(AsyncStreamHandler):
                 # Handle function calls
                 elif event.type == "response.function_call_arguments.start":
                     self.function_call_in_progress = True
                     self.current_function_args = ""
-                    self.call_id = event.call_id if hasattr(event, 'call_id') else None
                 elif event.type == "response.function_call_arguments.delta":
                     if self.function_call_in_progress:
@@ -784,6 +780,7 @@ class OpenAIHandler(AsyncStreamHandler):
                 elif event.type == "response.function_call_arguments.done":
                     if self.function_call_in_progress:
                         try:
                             args = json.loads(self.current_function_args)
                             query = args.get("query", "")
@@ -791,19 +788,19 @@ class OpenAIHandler(AsyncStreamHandler):
                             # Emit search event to client
                             await self.output_queue.put(AdditionalOutputs({
                                 "type": "search",
-                                "query": query,
-                                "results": True
                             }))
                             # Perform the search
                             search_results = await self.search_web(query)
                             # Send function result back to the model
-                            if self.connection and self.call_id:
                                 await self.connection.conversation.item.create(
                                     item={
                                         "type": "function_call_output",
-                                        "call_id": self.call_id,
                                         "output": search_results
                                     }
                                 )
@@ -814,7 +811,7 @@ class OpenAIHandler(AsyncStreamHandler):
                         finally:
                             self.function_call_in_progress = False
                             self.current_function_args = ""
-                            self.call_id = None
     async def receive(self, frame: tuple[int, np.ndarray]) -> None:
         if not self.connection:
@@ -833,80 +830,77 @@ class OpenAIHandler(AsyncStreamHandler):
             self.connection = None
-# Store connection settings
-connection_settings = {}
 # Initialize search client
 brave_api_key = os.getenv("BSEARCH_API")
 search_client = BraveSearchClient(brave_api_key) if brave_api_key else None
 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 handler
-handler = OpenAIHandler()
-handler.search_client = search_client
-# Create stream
-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,
-)
 app = FastAPI()
-stream.mount(app)
-@app.post("/update-search-setting")
-async def update_search_setting(request: dict):
-    """Update web search setting for a connection"""
     webrtc_id = request.get("webrtc_id")
     web_search_enabled = request.get("web_search_enabled", False)
-    if webrtc_id:
-        connection_settings[webrtc_id] = {"web_search_enabled": web_search_enabled}
-        # Update handler setting
-        handler.web_search_enabled = web_search_enabled
-        # Reconnect with updated settings if needed
-        if handler.connection:
-            await handler.shutdown()
-            await handler.start_up()
-    return {"status": "updated"}
-# Override the stream's offer handler
-original_offer = stream.offer
-async def custom_offer(request: dict):
-    """Handle WebRTC offer with web search preference"""
-    webrtc_id = request.get("webrtc_id")
-    web_search_enabled = request.get("web_search_enabled", False)
-    if webrtc_id:
-        connection_settings[webrtc_id] = {"web_search_enabled": web_search_enabled}
-        handler.web_search_enabled = web_search_enabled
-    return await original_offer(request)
-# Replace the offer method
-stream.offer = custom_offer
 @app.get("/outputs")
 async def outputs(webrtc_id: str):
     """Stream outputs including search events"""
     async def output_stream():
-        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':
@@ -932,8 +926,38 @@ if __name__ == "__main__":
     mode = os.getenv("MODE")
     if mode == "UI":
         stream.ui.launch(server_port=7860)
     elif mode == "PHONE":
         stream.fastphone(host="0.0.0.0", port=7860)
     else:
         uvicorn.run(app, host="0.0.0.0", port=7860)

         searchToggle.addEventListener('click', () => {
             webSearchEnabled = !webSearchEnabled;
             searchToggle.classList.toggle('active', webSearchEnabled);
+            console.log('Web search enabled:', webSearchEnabled);
         });
         function updateStatus(state) {
                 });
                 eventSource.addEventListener("search", (event) => {
                     const eventJson = JSON.parse(event.data);
+                    if (eventJson.query) {
                         addMessage("search-result", `웹 검색 중: "${eventJson.query}"`);
                     }
                 });
 class OpenAIHandler(AsyncStreamHandler):
+    def __init__(self, web_search_enabled: bool = False, search_client: Optional[BraveSearchClient] = None, webrtc_id: str = None) -> None:
         super().__init__(
             expected_layout="mono",
             output_sample_rate=SAMPLE_RATE,
         )
         self.connection = None
         self.output_queue = asyncio.Queue()
+        self.web_search_enabled = web_search_enabled
+        self.search_client = search_client
         self.function_call_in_progress = False
         self.current_function_args = ""
+        self.current_call_id = None
+        self.webrtc_id = webrtc_id
     def copy(self):
+        # Return a new instance with the same settings
+        return OpenAIHandler(self.web_search_enabled, self.search_client, self.webrtc_id)
     async def search_web(self, query: str) -> str:
         """Perform web search and return formatted results"""
         if not self.search_client or not self.web_search_enabled:
             return "웹 검색이 비활성화되어 있습니다."
+        print(f"Searching web for: {query}")
         results = await self.search_client.search(query)
         if not results:
             return f"'{query}'에 대한 검색 결과를 찾을 수 없습니다."
     async def start_up(self):
         """Connect to realtime API with function calling enabled"""
+        print(f"Starting up handler with web_search_enabled={self.web_search_enabled}")
         self.client = openai.AsyncOpenAI()
         # Define the web search function
                     }
                 }
             }]
+            print("Web search function added to tools")
         async with self.client.beta.realtime.connect(
             model="gpt-4o-mini-realtime-preview-2024-12-17"
             await conn.session.update(session=session_update)
             self.connection = conn
+            print(f"Connected with tools: {len(tools)} functions")
             async for event in self.connection:
                 if event.type == "response.audio_transcript.done":
                 # Handle function calls
                 elif event.type == "response.function_call_arguments.start":
+                    print(f"Function call started, call_id: {getattr(event, 'call_id', 'unknown')}")
                     self.function_call_in_progress = True
                     self.current_function_args = ""
+                    self.current_call_id = getattr(event, 'call_id', None)
                 elif event.type == "response.function_call_arguments.delta":
                     if self.function_call_in_progress:
                 elif event.type == "response.function_call_arguments.done":
                     if self.function_call_in_progress:
+                        print(f"Function call done, args: {self.current_function_args}")
                         try:
                             args = json.loads(self.current_function_args)
                             query = args.get("query", "")
                             # Emit search event to client
                             await self.output_queue.put(AdditionalOutputs({
                                 "type": "search",
+                                "query": query
                             }))
                             # Perform the search
                             search_results = await self.search_web(query)
+                            print(f"Search results length: {len(search_results)}")
                             # Send function result back to the model
+                            if self.connection and self.current_call_id:
                                 await self.connection.conversation.item.create(
                                     item={
                                         "type": "function_call_output",
+                                        "call_id": self.current_call_id,
                                         "output": search_results
                                     }
                                 )
                         finally:
                             self.function_call_in_progress = False
                             self.current_function_args = ""
+                            self.current_call_id = None
     async def receive(self, frame: tuple[int, np.ndarray]) -> None:
         if not self.connection:
             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
+print(f"Search client initialized: {search_client is not None}, API key present: {bool(brave_api_key)}")
 def update_chatbot(chatbot: list[dict], response: ResponseAudioTranscriptDoneEvent):
     chatbot.append({"role": "assistant", "content": response.transcript})
     return chatbot
+# Create base stream without handler for mounting
 app = FastAPI()
+@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,
+    )
+    # Mount stream temporarily
+    stream.mount(app)
+    # Store stream reference
+    handler.stream = stream
+    # Process the offer
+    response = await stream.offer(request)
+    return response
 @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):
             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':
     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)