Update app.py

app.py

@@ -1,34 +1,50 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-from transformers import AutoTokenizer
-from langchain.memory import ConversationBufferWindowMemory
-from langchain.schema import HumanMessage, AIMessage, SystemMessage
+from transformers import AutoTokenizer  # Import the tokenizer
+from langchain.memory import ConversationBufferMemory
+from langchain.schema import HumanMessage, AIMessage
 
-# Initialize tokenizer and inference client
+# Use the appropriate tokenizer for your model.
 tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
-MAX_CONTEXT_LENGTH = 4096
+# Define a maximum context length (tokens).  Check your model's documentation!
+MAX_CONTEXT_LENGTH = 4096  # Example: Adjust this based on your model!
 
-# Load prompt from file
+# Read the default prompt from a file
 with open("prompt.txt", "r") as file:
     nvc_prompt_template = file.read()
 
-# Initialize LangChain Memory (buffer window to keep recent conversation)
-memory = ConversationBufferWindowMemory(k=10, return_messages=True)
+# Initialize LangChain Conversation Memory
+memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
 def count_tokens(text: str) -> int:
+    """Counts the number of tokens in a given string."""
     return len(tokenizer.encode(text))
 
-def truncate_history(messages, max_length):
+def truncate_memory(memory, system_message: str, max_length: int):
+    """
+    Truncates the conversation memory messages to fit within the maximum token limit.
+    
+    Args:
+        memory: The LangChain conversation memory object.
+        system_message: The system message.
+        max_length: The maximum number of tokens allowed.
+        
+    Returns:
+        A list of messages (as dicts with role and content) that fit within the token limit.
+    """
     truncated_messages = []
-    total_tokens = 0
-
-    for message in reversed(messages):
-        message_tokens = count_tokens(message.content)
-        if total_tokens + message_tokens <= max_length:
-            truncated_messages.insert(0, message)
-            total_tokens += message_tokens
+    system_tokens = count_tokens(system_message)
+    current_length = system_tokens
+
+    # Iterate backwards through the memory (newest to oldest)
+    for msg in reversed(memory.chat_memory.messages):
+        tokens = count_tokens(msg.content)
+        if current_length + tokens <= max_length:
+            role = "user" if isinstance(msg, HumanMessage) else "assistant"
+            truncated_messages.insert(0, {"role": role, "content": msg.content})
+            current_length += tokens
         else:
             break
 
@@ -36,38 +52,39 @@ def truncate_history(messages, max_length):
 
 def respond(
     message,
-    history,
+    history: list[tuple[str, str]],  # Required by Gradio but we now use LangChain memory
     system_message,
     max_tokens,
     temperature,
     top_p,
 ):
+    """
+    Responds to a user message while maintaining conversation history via LangChain memory.
+    It builds the prompt with a system message and the (truncated) conversation history,
+    streams the response from the client, and finally updates the memory with the new response.
+    """
+    # Use your prompt template as the system message.
     formatted_system_message = nvc_prompt_template
 
-    # Retrieve conversation history from LangChain memory
-    memory.save_context({"input": message}, {"output": ""})
-    chat_history = memory.load_memory_variables({})["history"]
-
-    # Truncate history to ensure it fits within context window
-    max_history_tokens = MAX_CONTEXT_LENGTH - max_tokens - count_tokens(formatted_system_message) - 100
-    truncated_chat_history = truncate_history(chat_history, max_history_tokens)
+    # Prepare and add the new user message (with your special tokens) to memory.
+    new_user_message = f"<|user|>\n{message}</s>"
+    memory.chat_memory.add_message(HumanMessage(content=new_user_message))
 
-    # Construct the messages for inference
-    messages = [SystemMessage(content=formatted_system_message)]
-    messages.extend(truncated_chat_history)
-    messages.append(HumanMessage(content=message))
+    # Truncate memory to ensure the context fits within the maximum token length (reserve space for generation).
+    truncated_history = truncate_memory(
+        memory, formatted_system_message, MAX_CONTEXT_LENGTH - max_tokens - 100
+    )
+    # Ensure the current user message is present at the end.
+    if not truncated_history or truncated_history[-1]["content"] != new_user_message:
+        truncated_history.append({"role": "user", "content": new_user_message})
 
-    # Convert LangChain messages to the format required by HuggingFace client
-    formatted_messages = []
-    for msg in messages:
-        role = "system" if isinstance(msg, SystemMessage) else "user" if isinstance(msg, HumanMessage) else "assistant"
-        content = f"<|{role}|>\n{msg.content}</s>"
-        formatted_messages.append({"role": role, "content": content})
+    # Build the full message list: system prompt + conversation history.
+    messages = [{"role": "system", "content": formatted_system_message}] + truncated_history
 
     response = ""
     try:
         for chunk in client.chat_completion(
-            formatted_messages,
+            messages,
             max_tokens=max_tokens,
             stream=True,
             temperature=temperature,
@@ -76,14 +93,13 @@ def respond(
             token = chunk.choices[0].delta.content
             response += token
             yield response
-
-        # Save AI's response in LangChain memory
-        memory.chat_memory.add_ai_message(response)
-
     except Exception as e:
         print(f"An error occurred: {e}")
         yield "I'm sorry, I encountered an error. Please try again."
 
+    # Once the full response is generated, add it to the LangChain memory.
+    memory.chat_memory.add_message(AIMessage(content=f"<|assistant|>\n{response}</s>"))
+
 # --- Gradio Interface ---
 demo = gr.ChatInterface(
     respond,