Update app.py

app.py

@@ -4,42 +4,24 @@ from huggingface_hub import InferenceClient
 # Initialize the InferenceClient with the model ID from Hugging Face
 client = InferenceClient(model="HuggingFaceH4/zephyr-7b-beta")
 
-def respond(
-    message: str,
-    history: list[tuple[str, str]],
-    system_message: str,
+def generate_response(
+    messages: list[dict],
     max_tokens: int,
     temperature: float,
     top_p: float,
 ):
     """
-    Generates a response from the AI model based on the user's message and chat history.
-    
+    Generates a response from the AI model based on the provided messages.
+
     Args:
-        message (str): The user's input message.
-        history (list): A list of tuples representing the conversation history (user, assistant).
-        system_message (str): A system-level message guiding the AI's behavior.
+        messages (list): A list of messages representing the conversation history.
         max_tokens (int): The maximum number of tokens for the output.
-        temperature (float): Sampling temperature for controlling the randomness.
+        temperature (float): Sampling temperature for controlling randomness.
         top_p (float): Top-p (nucleus sampling) for controlling diversity.
 
     Yields:
         str: The AI's response as it is generated.
     """
-    
-    # Prepare the conversation history for the API call
-    messages = [{"role": "system", "content": system_message}]
-    
-    for user_input, assistant_response in history:
-        if user_input:
-            messages.append({"role": "user", "content": user_input})
-        if assistant_response:
-            messages.append({"role": "assistant", "content": assistant_response})
-
-    # Add the latest user message to the conversation
-    messages.append({"role": "user", "content": message})
-
-    # Initialize an empty response
     response = ""
 
     try:
@@ -58,6 +40,56 @@ def respond(
     except Exception as e:
         yield f"An error occurred: {str(e)}"
 
+def build_messages(system_message: str, history: list[tuple[str, str]], user_message: str) -> list[dict]:
+    """
+    Builds the list of messages for the model based on system message, history, and latest user message.
+
+    Args:
+        system_message (str): A system-level message guiding the AI's behavior.
+        history (list): A list of tuples representing the conversation history (user, assistant).
+        user_message (str): The latest user message.
+
+    Returns:
+        list: A list of message dictionaries formatted for the API call.
+    """
+    messages = [{"role": "system", "content": system_message}]
+    
+    for user_input, assistant_response in history:
+        if user_input:
+            messages.append({"role": "user", "content": user_input})
+        if assistant_response:
+            messages.append({"role": "assistant", "content": assistant_response})
+
+    # Add the latest user message to the conversation
+    messages.append({"role": "user", "content": user_message})
+    
+    return messages
+
+def respond(
+    message: str,
+    history: list[tuple[str, str]],
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+):
+    """
+    Handles the interaction with the model to generate a response based on user input and chat history.
+
+    Args:
+        message (str): The user's input message.
+        history (list): A list of tuples representing the conversation history (user, assistant).
+        system_message (str): A system-level message guiding the AI's behavior.
+        max_tokens (int): The maximum number of tokens for the output.
+        temperature (float): Sampling temperature for controlling randomness.
+        top_p (float): Top-p (nucleus sampling) for controlling diversity.
+
+    Yields:
+        str: The AI's response as it is generated.
+    """
+    messages = build_messages(system_message, history, message)
+    yield from generate_response(messages, max_tokens, temperature, top_p)
+
 # Define the UI layout with improved structure
 with gr.Blocks() as demo:
     gr.Markdown("# 🧠 AI Chatbot Interface")
@@ -96,19 +128,21 @@ with gr.Blocks() as demo:
             clear_btn = gr.Button("Clear Chat")
 
     # Handle sample prompt selection
-    def update_message(prompt):
+    def update_message(prompt: str) -> str:
         return prompt
 
     sample_prompt.change(fn=update_message, inputs=sample_prompt, outputs=message)
     
     # Update the chatbot with the new message and response
-    submit_btn.click(respond, 
-                     inputs=[message, chatbot, system_message, max_tokens, temperature, top_p], 
-                     outputs=[chatbot],
-                     show_progress=True)
+    submit_btn.click(
+        fn=respond, 
+        inputs=[message, chatbot, system_message, max_tokens, temperature, top_p], 
+        outputs=[chatbot],
+        show_progress=True
+    )
     
     # Clear the chat history
-    def clear_chat():
+    def clear_chat() -> list:
         return []
 
     clear_btn.click(fn=clear_chat, inputs=None, outputs=chatbot)