Update app.py

app.py

@@ -2,7 +2,6 @@ import streamlit as st
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain.schema import HumanMessage, AIMessage
-from langchain.chat_models.base import BaseChatModel
 from llama_cpp import Llama
 
 # Initialize the Llama model
@@ -15,28 +14,11 @@ llm = Llama.from_pretrained(
     chat_format="chatml"
 )
 
-# Define the LangChain model for Llama
-class LlamaChatModel(BaseChatModel):
-    def _llm_type(self) -> str:
-        return "llama"
-
-    def _generate(self, messages, stop=None):
-        # Prepare prompt from conversation history
-        prompt = "\n".join(
-            f"user: {msg.content}" if isinstance(msg, HumanMessage) else f"assistant: {msg.content}"
-            for msg in messages
-        )
-
-        # Generate response from Llama
-        response = llm.chat(prompt)
-        return [AIMessage(content=response)]
-
-# Initialize memory and chat model
+# Initialize memory
 memory = ConversationBufferMemory(return_messages=True)
-llama_chat_model = LlamaChatModel()
 
-# Create the conversation chain
-conversation = ConversationChain(memory=memory, llm=llama_chat_model)
+# Create the conversation chain directly using the Llama model
+conversation = ConversationChain(memory=memory, llm=llm)
 
 # Streamlit UI
 st.title("Chatbot with LangChain and Llama")