Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import uuid
 from langchain_huggingface import HuggingFaceEndpoint
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.messages import AIMessage
@@ -8,13 +9,12 @@ from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 import gradio as gr
 
-
 # Set your API keys from environment variables
 langchain_key = os.getenv("LANGCHAIN_API_KEY")
 HF_key = os.getenv("HUGGINGFACEHUB_TOKEN")
-LANGCHAIN_TRACING_V2=True
-LANGCHAIN_ENDPOINT="https://api.smith.langchain.com"
-LANGCHAIN_PROJECT="LLM_CHATBOT"
+LANGCHAIN_TRACING_V2 = True
+LANGCHAIN_ENDPOINT = "https://api.smith.langchain.com"
+LANGCHAIN_PROJECT = "LLM_CHATBOT"
 
 os.environ["LANGCHAIN_TRACING_V2"] = str(LANGCHAIN_TRACING_V2)
 os.environ["LANGCHAIN_API_KEY"] = langchain_key
@@ -28,7 +28,7 @@ llm = HuggingFaceEndpoint(
     task="text-generation",
     max_new_tokens=150,
     do_sample=False,
-    token =HF_key
+    token=HF_key
 )
 
 # Create a Chat Prompt Template
@@ -56,18 +56,33 @@ with_message_history = RunnableWithMessageHistory(chain, get_session_history)
 def chat(session_id, user_input):
     config = {"configurable": {"session_id": session_id}}
     human_message = HumanMessage(content=user_input)
-    response = with_message_history.invoke({"messages": [human_message]}, config=config)
+    try:
+        response = with_message_history.invoke({"messages": [human_message]}, config=config)
+    except Exception as e:
+        response = AIMessage(content=f"An error occurred: {str(e)}")
     return response
 
+# Gradio session function
+def get_session_id(request: gr.Request):
+    session_id = request.cookies.get("session_id")
+    if not session_id:
+        session_id = str(uuid.uuid4())
+        response = gr.Response().cookies.set("session_id", session_id)
+    else:
+        response = gr.Response()
+    return session_id, response
+
 # Gradio interface
 iface = gr.Interface(
     fn=chat,
-    inputs=[gr.Textbox(lines=1, placeholder="Enter Session ID"), gr.Textbox(lines=7, placeholder="Enter your message")],
+    inputs=[gr.Textbox(lines=1, placeholder="Enter your message")],
     outputs="text",
     title="LangChain Chatbot",
-    description="A chatbot that remembers your past interactions. Enter your session ID and message."
+    description="A chatbot that remembers your past interactions. Enter your message."
 )
 
+# Add session management
+iface = iface.wrap_inputs(get_session_id, gr.Textbox)
+
 # Launch the app
 iface.launch()
-