Update app.py

app.py

@@ -2,6 +2,7 @@ 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
@@ -14,11 +15,28 @@ llm = Llama.from_pretrained(
     chat_format="chatml"
 )
 
-# Initialize memory
+# 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
 memory = ConversationBufferMemory(return_messages=True)
+llama_chat_model = LlamaChatModel()
 
-# Create the conversation chain directly using the Llama model
-conversation = ConversationChain(memory=memory, llm=llm)
+# Create the conversation chain
+conversation = ConversationChain(memory=memory, llm=llama_chat_model)
 
 # Streamlit UI
 st.title("Chatbot with LangChain and Llama")