update: add streamimng

app.py

@@ -411,11 +411,21 @@ Rank topics based on their potential virality and engagement for social media cl
                 },
                 {"role": "user", "content": prompt},
             ],
+            stream=True,
         )
-        return completion.choices[0].message.content
+
+        collected_messages = []
+        # Iterate through the stream
+        for chunk in completion:
+            if chunk.choices[0].delta.content is not None:
+                chunk_message = chunk.choices[0].delta.content
+                collected_messages.append(chunk_message)
+                # Yield the accumulated message so far
+                yield "".join(collected_messages)
+
     except Exception as e:
         print(f"Error in initial analysis: {str(e)}")
-        return "An error occurred during initial analysis. Please check your API key and file path."
+        yield "An error occurred during initial analysis. Please check your API key and file path."
 
 
 def chat(
@@ -498,63 +508,77 @@ In the URL, make sure that after RSID there is ? and then rest of the fields are
         messages.append({"role": "user", "content": message})
 
         completion = client.chat.completions.create(
-            model="gpt-4o-mini",
-            messages=messages,
-            tools=tools,
+            model="gpt-4o-mini", messages=messages, tools=tools, stream=True
         )
+        collected_messages = []
+        tool_calls_detected = False
+
+        for chunk in completion:
+            if chunk.choices[0].delta.tool_calls:
+                tool_calls_detected = True
+                # Handle tool calls without streaming
+                response = client.chat.completions.create(
+                    model="gpt-4o-mini",
+                    messages=messages,
+                    tools=tools,
+                )
 
-        response = completion.choices[0].message
-        messages.append(response)
-
-        if response.tool_calls:
-            if response.tool_calls[0].function.name == "correct_speaker_name_with_url":
-                args = eval(response.tool_calls[0].function.arguments)
-                url = args.get("url", None)
-                if url:
-                    transcript_processor.correct_speaker_mapping_with_agenda(url)
-                    corrected_speaker_mapping = transcript_processor.speaker_mapping
-                    function_call_result_message = {
-                        "role": "tool",
-                        "content": json.dumps(
-                            {
-                                "speaker_mapping": f"Corrected Speaker Mapping is: {corrected_speaker_mapping}\n, All speakers should be addressed via this mapping. The next message should be comparing old speaker names (do not use spk_0, spk_1, use the old names) and corrected speaker names.",
+                if response.choices[0].message.tool_calls:
+                    tool_call = response.choices[0].message.tool_calls[0]
+                    if tool_call.function.name == "correct_speaker_name_with_url":
+                        args = eval(tool_call.function.arguments)
+                        url = args.get("url", None)
+                        if url:
+                            transcript_processor.correct_speaker_mapping_with_agenda(
+                                url
+                            )
+                            corrected_speaker_mapping = (
+                                transcript_processor.speaker_mapping
+                            )
+                            function_call_result_message = {
+                                "role": "tool",
+                                "content": json.dumps(
+                                    {"speaker_mapping": f"Corrected Speaker Mapping..."}
+                                ),
+                                "name": tool_call.function.name,
+                                "tool_call_id": tool_call.id,
                             }
-                        ),
-                        "name": response.tool_calls[0].function.name,
-                        "tool_call_id": response.tool_calls[0].id,
-                    }
-
-                    # messages.append(response.choices[0]["message"])
-                    messages.append(function_call_result_message)
-                    completion_payload = {"model": "gpt-4o-mini", "messages": messages}
-                    # print("messages", messages[3])
-                    response = client.chat.completions.create(**completion_payload)
-                    # print("no error here")
-
-                    return response.choices[0].message.content
-
-                else:
-                    return "No URL provided for correcting speaker names."
-            elif response.tool_calls[0].function.name == "correct_call_type":
-                args = eval(response.tool_calls[0].function.arguments)
-                call_type = args.get("call_type", None)
-                if call_type:
-                    analysis = get_initial_analysis(
-                        transcript_processor, call_type, rsid, origin, call_type
-                    )
-                    return analysis
-                else:
-                    return "No call type provided for correction."
-            else:
-                return "No tool for this request"
-        return response.content
+                            messages.append(function_call_result_message)
+
+                            # Get final response after tool call
+                            final_response = client.chat.completions.create(
+                                model="gpt-4o-mini", messages=messages, stream=True
+                            )
+
+                            # Stream the final response
+                            for final_chunk in final_response:
+                                if final_chunk.choices[0].delta.content:
+                                    yield final_chunk.choices[0].delta.content
+                            return
+
+                    elif tool_call.function.name == "correct_call_type":
+                        args = eval(tool_call.function.arguments)
+                        call_type = args.get("call_type", None)
+                        if call_type:
+                            # Stream the analysis for corrected call type
+                            for content in get_initial_analysis(
+                                transcript_processor, call_type, rsid, origin, call_type
+                            ):
+                                yield content
+                            return
+                break  # Exit streaming loop if tool calls detected
+
+            if not tool_calls_detected and chunk.choices[0].delta.content is not None:
+                chunk_message = chunk.choices[0].delta.content
+                collected_messages.append(chunk_message)
+                yield "".join(collected_messages)
 
     except Exception as e:
         print(f"Unexpected error in chat: {str(e)}")
         import traceback
 
         print(f"Traceback: {traceback.format_exc()}")
-        return "Sorry, there was an error processing your request."
+        yield "Sorry, there was an error processing your request."
 
 
 def create_chat_interface():
@@ -688,18 +712,62 @@ def create_chat_interface():
         iframe_html = "<iframe id='link-frame'></iframe>"
         gr.HTML(value=iframe_html)  # Add iframe to the UI
 
+        def respond(
+            message: str,
+            chat_history: List,
+            transcript_processor,
+            cid,
+            rsid,
+            origin,
+            ct,
+        ):
+            if not transcript_processor:
+                bot_message = "Transcript processor not initialized."
+                chat_history.append((message, bot_message))
+                return "", chat_history
+
+            chat_history.append((message, ""))
+            for chunk in chat(
+                message,
+                chat_history[:-1],  # Exclude the current incomplete message
+                transcript_processor,
+                cid,
+                rsid,
+                origin,
+                ct,
+            ):
+                chat_history[-1] = (message, chunk)
+                yield "", chat_history
+
+        msg.submit(
+            respond,
+            [
+                msg,
+                chatbot,
+                transcript_processor_state,
+                call_id_state,
+                colab_id_state,
+                origin_state,
+                ct_state,
+            ],
+            [msg, chatbot],
+        )
+
+        # Handle initial loading with streaming
         def on_app_load(request: gr.Request):
             cid = request.query_params.get("cid", None)
             rsid = request.query_params.get("rsid", None)
             origin = request.query_params.get("origin", None)
             ct = request.query_params.get("ct", None)
             turl = request.query_params.get("turl", None)
+
             required_params = ["cid", "rsid", "origin", "ct", "turl"]
             missing_params = [
                 param
                 for param in required_params
                 if request.query_params.get(param) is None
             ]
+
             if missing_params:
                 error_message = (
                     f"Missing required parameters: {', '.join(missing_params)}"
@@ -715,33 +783,16 @@ def create_chat_interface():
                     None,
                 ]
 
-            # if any param is missing, return error
-            if not cid or not rsid or not origin or not ct or not turl:
-                error_message = "Error processing"
-                chatbot_value = [(None, error_message)]
-                return [
-                    chatbot_value,
-                    None,
-                    None,
-                    None,
-                    None,
-                    None,
-                    None,
-                ]
-
             try:
                 transcript_data = get_transcript_for_url(turl)
                 transcript_processor = TranscriptProcessor(
                     transcript_data=transcript_data
                 )
-                initial_analysis = get_initial_analysis(
-                    transcript_processor, cid, rsid, origin, ct
-                )
 
-                chatbot_value = [
-                    (None, initial_analysis)
-                ]  # initialized with initial analysis and assistant is None
+                # Initialize with empty message
+                chatbot_value = [(None, "")]
 
+                # Return initial values with the transcript processor
                 return [
                     chatbot_value,
                     transcript_processor,
@@ -764,6 +815,19 @@ def create_chat_interface():
                     None,
                 ]
 
+        def stream_initial_analysis(
+            chatbot_value, transcript_processor, cid, rsid, origin, ct
+        ):
+            if transcript_processor:
+                for chunk in get_initial_analysis(
+                    transcript_processor, cid, rsid, origin, ct
+                ):
+                    chatbot_value[0] = (None, chunk)
+                    yield chatbot_value
+            else:
+                yield chatbot_value
+
+        # Modified load event to handle streaming
         demo.load(
             on_app_load,
             inputs=None,
@@ -776,36 +840,9 @@ def create_chat_interface():
                 ct_state,
                 turl_state,
             ],
-        )
-
-        def respond(
-            message: str,
-            chat_history: List,
-            transcript_processor,
-            cid,
-            rsid,
-            origin,
-            ct,
-        ):
-            if not transcript_processor:
-                bot_message = "Transcript processor not initialized."
-            else:
-                bot_message = chat(
-                    message,
-                    chat_history,
-                    transcript_processor,
-                    cid,
-                    rsid,
-                    origin,
-                    ct,
-                )
-            chat_history.append((message, bot_message))
-            return "", chat_history
-
-        msg.submit(
-            respond,
-            [
-                msg,
+        ).then(
+            stream_initial_analysis,
+            inputs=[
                 chatbot,
                 transcript_processor_state,
                 call_id_state,
@@ -813,7 +850,7 @@ def create_chat_interface():
                 origin_state,
                 ct_state,
             ],
-            [msg, chatbot],
+            outputs=[chatbot],
         )
 
     return demo