Update app.py

app.py

@@ -7,25 +7,31 @@ MODEL = "nomiChroma3.1"
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 def respond(
-    message,
-    history: list[tuple[str, str]],
-    max_tokens,
-    temperature,
-    top_p,
-):
-    # Define system message internally instead of as a user input
+    message: str,
+    chat_history: list[tuple[str, str]],
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+) -> tuple[list[tuple[str, str]], str]:
+    """
+    Generate a response and update chat history.
+    Returns tuple of (new_history, None) to clear input box.
+    """
     system_message = "You are a maritime legal assistant with expertise strictly in Indian maritime law. Provide detailed legal advice and information based on Indian maritime legal principles and regulations."
     
     messages = [{"role": "system", "content": system_message}]
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    for user_msg, assistant_msg in chat_history:
+        messages.extend([
+            {"role": "user", "content": user_msg},
+            {"role": "assistant", "content": assistant_msg}
+        ])
     messages.append({"role": "user", "content": message})
+    
+    chat_history = chat_history + [(message, None)]
     response = ""
+    
     try:
-        for message in client.chat_completion(
+        for chunk in client.chat_completion(
             messages,
             max_tokens=max_tokens,
             stream=True,
@@ -33,34 +39,41 @@ def respond(
             top_p=top_p,
         ):
             try:
-                if isinstance(message, ChatCompletionStreamOutput):
-                    content = message.choices[0].delta.content
-                    if content is not None:
+                if isinstance(chunk, ChatCompletionStreamOutput):
+                    content = chunk.choices[0].delta.content
+                    if content:
                         response += content
-                        yield response
-                    if message.choices[0].finish_reason == 'stop':
+                        chat_history[-1] = (message, response)
+                        yield chat_history, ""
+                    if chunk.choices[0].finish_reason == 'stop':
                         break
-                elif isinstance(message, dict):
-                    content = message.get('choices', [{}])[0].get('delta', {}).get('content')
+                elif isinstance(chunk, dict):
+                    content = chunk.get('choices', [{}])[0].get('delta', {}).get('content')
                     if content:
                         response += content
-                        yield response
-                    if message.get('choices', [{}])[0].get('finish_reason') == 'stop':
+                        chat_history[-1] = (message, response)
+                        yield chat_history, ""
+                    if chunk.get('choices', [{}])[0].get('finish_reason') == 'stop':
                         break
-                elif isinstance(message, str):
-                    if message.strip():
-                        response += message
-                        yield response
+                elif isinstance(chunk, str) and chunk.strip():
+                    response += chunk
+                    chat_history[-1] = (message, response)
+                    yield chat_history, ""
+                    
             except Exception as e:
-                print(f"Error processing message: {e}")
+                print(f"Error processing chunk: {e}")
                 continue
 
-        if response:
-            yield response
+        if not response:
+            chat_history[-1] = (message, "I apologize, but I couldn't generate a response. Please try again.")
+        
+        yield chat_history, ""
 
     except Exception as e:
-        print(f"An error occurred: {e}")
-        yield f"An error occurred: {e}"
+        error_msg = f"An error occurred: {str(e)}"
+        chat_history[-1] = (message, error_msg)
+        yield chat_history, ""
+
 
 custom_css = """
     /* Global styles */
@@ -183,13 +196,11 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:
         </div>
     """)
 
-    # Main layout with sidebar
     with gr.Row():
         # Sidebar
         with gr.Column(scale=1, elem_classes="sidebar"):
             gr.Markdown("### Example Queries", elem_classes="sidebar-title")
             
-            # Example queries in sidebar
             example_queries = [
                 "What are the key regulations governing ports in India?",
                 "Explain the concept of cabotage in Indian maritime law.",
@@ -199,7 +210,6 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:
                 "What are the rules for coastal cargo transportation in India?"
             ]
             
-            # Create buttons for example queries
             with gr.Column(elem_classes="example-queries"):
                 example_buttons = [gr.Button(query, elem_classes="example-query-button") for query in example_queries]
             
@@ -239,29 +249,30 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:
                 clear = gr.Button("Clear")
 
             # Event handlers
-            submit_click = submit.click(
+            msg.submit(
                 respond,
                 inputs=[msg, chatbot, max_tokens, temperature, top_p],
-                outputs=chatbot
+                outputs=[chatbot, msg]
             )
-            msg_submit = msg.submit(
+            
+            submit.click(
                 respond,
                 inputs=[msg, chatbot, max_tokens, temperature, top_p],
-                outputs=chatbot
+                outputs=[chatbot, msg]
             )
-            clear.click(lambda: None, None, chatbot, queue=False)
             
-            # Add click handlers for example queries
+            clear.click(lambda: ([], ""), None, [chatbot, msg], queue=False)
+            
+            # Example query handlers
             for button in example_buttons:
                 button.click(
-                    lambda query: query,
+                    lambda query: (query, []),  # Reset chat history when using example
                     inputs=[button],
-                    outputs=[msg]
+                    outputs=[msg, chatbot],
                 ).then(
                     respond,
                     inputs=[msg, chatbot, max_tokens, temperature, top_p],
-                    outputs=chatbot
-                )
+                    outputs=[chatbot, msg]
 
 if __name__ == "__main__":
     demo.launch()