Update app.py

app.py

@@ -78,7 +78,7 @@ def get_model_response(messages, inner_messages, message, system_message):
     messages_for_model.extend(inner_messages)
 
     try:
-        return client.chat.completions.create(
+        response = client.chat.completions.create(
             model=MODEL,
             messages=messages_for_model,
             tools=tools,
@@ -86,6 +86,7 @@ def get_model_response(messages, inner_messages, message, system_message):
             top_p=0.65,
             max_tokens=4096,
         )
+        return response.choices[0].message
     except Exception as e:
         print(f"An error occurred while getting model response: {str(e)}")
         print(messages_for_model)
@@ -93,7 +94,6 @@ def get_model_response(messages, inner_messages, message, system_message):
 
 def respond(message, history, system_message):
     inner_history = []
-
     available_functions = {
         "legal_tool_function": legal_tool_function,
     }
@@ -102,13 +102,14 @@ def respond(message, history, system_message):
     assistant_native_message_list = []
 
     while True:
-        response_message = (
-            get_model_response(history, inner_history, message, system_message)
-            .choices[0]
-            .message
-        )
+        response_message = get_model_response(history, inner_history, message, system_message)
+
+        if response_message is None:
+            return assistant_content, history
 
         if not response_message.tool_calls and response_message.content is not None:
+            assistant_content += response_message.content
+            assistant_native_message_list.append(response_message)
             break
 
         if response_message.tool_calls is not None:
@@ -132,7 +133,8 @@ def respond(message, history, system_message):
                 "metadata": {"native_messages": assistant_native_message_list},
             }
 
-            yield assistant_message
+            # Collect responses
+            response_list = [assistant_message]
 
             for tool_call in response_message.tool_calls:
                 function_response = call_function(tool_call, available_functions)
@@ -161,19 +163,10 @@ def respond(message, history, system_message):
                     "content": assistant_content,
                     "metadata": {"native_messages": assistant_native_message_list},
                 }
-                yield tool_message
+                response_list.append(tool_message)
                 inner_history.append(native_tool_message)
 
-    assistant_content += response_message.content
-    assistant_native_message_list.append(response_message)
-
-    final_message = {
-        "role": "assistant",
-        "content": assistant_content,
-        "metadata": {"native_messages": assistant_native_message_list},
-    }
-    
-    yield final_message
+            return response_list, inner_history
 
 # Update the system prompt to be more relevant to maritime legal assistance
 system_prompt = "You are a maritime legal assistant with expertise in maritime law. Provide detailed legal advice and information based on maritime legal principles and regulations."
@@ -185,106 +178,10 @@ with gr.Blocks() as demo:
     message_input = gr.Textbox(label="Message")
     
     def process_message(message, history, system_message):
-        response_gen = respond(message, history, system_message)
-        return list(response_gen), history
+        responses, updated_history = respond(message, history, system_message)
+        return responses, updated_history
 
     message_input.submit(process_message, [message_input, chatbot, system_message_input], [chatbot, chatbot])
 
-if __name__ == "__main__":
-    demo.launch()
-
-    assistant_content = ""
-    assistant_native_message_list = []
-
-    while True:
-        response_message = (
-            get_model_response(history, inner_history, message, system_message)
-            .choices[0]
-            .message
-        )
-
-        if not response_message.tool_calls and response_message.content is not None:
-            break
-
-        if response_message.tool_calls is not None:
-            assistant_native_message_list.append(response_message)
-            inner_history.append(response_message)
-
-            assistant_content += (
-                "```json\n"
-                + json.dumps(
-                    [
-                        tool_call.model_dump()
-                        for tool_call in response_message.tool_calls
-                    ],
-                    indent=2,
-                )
-                + "\n```\n"
-            )
-            assistant_message = {
-                "role": "assistant",
-                "content": assistant_content,
-                "metadata": {"native_messages": assistant_native_message_list},
-            }
-
-            yield assistant_message
-
-            for tool_call in response_message.tool_calls:
-                function_response = call_function(tool_call, available_functions)
-                assistant_content += (
-                    "```json\n"
-                    + json.dumps(
-                        {
-                            "name": tool_call.function.name,
-                            "arguments": json.loads(tool_call.function.arguments),
-                            "response": json.loads(function_response["content"]),
-                        },
-                        indent=2,
-                    )
-                    + "\n```\n"
-                )
-                native_tool_message = {
-                    "tool_call_id": tool_call.id,
-                    "role": "tool",
-                    "content": function_response["content"],
-                }
-                assistant_native_message_list.append(
-                    native_tool_message
-                )
-                tool_message = {
-                    "role": "assistant",
-                    "content": assistant_content,
-                    "metadata": {"native_messages": assistant_native_message_list},
-                }
-                yield tool_message
-                inner_history.append(native_tool_message)
-
-    assistant_content += response_message.content
-    assistant_native_message_list.append(response_message)
-
-    final_message = {
-        "role": "assistant",
-        "content": assistant_content,
-        "metadata": {"native_messages": assistant_native_message_list},
-    }
-    
-    yield final_message
-
-# Update the system prompt to be more relevant to maritime legal assistance
-system_prompt = "You are a maritime legal assistant with expertise in maritime law. Provide detailed legal advice and information based on maritime legal principles and regulations."
-
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(
-            value=system_prompt,
-            label="System message",
-        ),
-    ],
-    type="messages",
-    title="Maritime Legal Assistant Chat",
-    description="This chatbot uses the fine-tuned maritime legal model to provide legal assistance and information related to maritime law.",
-)
-
 if __name__ == "__main__":
     demo.launch()