Create app.py

app.py

@@ -0,0 +1,158 @@
+import streamlit as st
+import os
+from langgraph.graph import MessagesState, StateGraph, START, END
+from typing_extensions import TypedDict
+from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage, AIMessage
+from typing import Annotated
+from langgraph.graph.message import add_messages
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_groq import ChatGroq
+
+# Define the state
+class MessagesState(TypedDict):
+    messages: Annotated[list[AnyMessage], add_messages]
+
+# Create graph function
+def create_chat_graph(system_prompt, model_name):
+    # Initialize LLM
+    llm = ChatGroq(model=model_name)
+    
+    # Create system message
+    system_message = SystemMessage(content=system_prompt)
+    
+    # Define the assistant function
+    def assistant(state: MessagesState):
+        # Get all messages including the system message
+        messages = [system_message] + state["messages"]
+        # Generate response
+        response = llm.invoke(messages)
+        # Return the response
+        return {"messages": [response]}
+    
+    # Initialize the graph builder
+    builder = StateGraph(MessagesState)
+    
+    # Add the assistant node
+    builder.add_node("assistant", assistant)
+    
+    # Define edges
+    builder.add_edge(START, "assistant")
+    builder.add_edge("assistant", END)
+    
+    # Create memory saver for persistence
+    memory = MemorySaver()
+    
+    # Compile the graph with memory
+    graph = builder.compile(checkpointer=memory)
+    
+    return graph
+
+# Set up Streamlit page
+st.set_page_config(page_title="Conversational AI Assistant", page_icon="💬")
+st.title("AI Chatbot with Memory")
+
+# Sidebar configuration
+st.sidebar.header("Configuration")
+
+# API Key input (using st.secrets in production)
+if "GROQ_API_KEY" not in os.environ:
+    api_key = st.sidebar.text_input("Enter your Groq API Key:", type="password")
+    if api_key:
+        os.environ["GROQ_API_KEY"] = api_key
+    else:
+        st.sidebar.warning("Please enter your Groq API key to continue.")
+
+# Model selection
+model_options = ["llama3-70b-8192", "mixtral-8x7b-32768", "gemma-7b-it"]
+selected_model = st.sidebar.selectbox("Select Model:", model_options)
+
+# System prompt
+default_prompt = "You are a helpful and friendly assistant. Maintain a conversational tone and remember previous interactions with the user."
+system_prompt = st.sidebar.text_area("System Prompt:", value=default_prompt, height=150)
+
+# Session ID for this conversation
+if "session_id" not in st.session_state:
+    import uuid
+    st.session_state.session_id = str(uuid.uuid4())
+
+# Initialize or get chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Initialize the graph on first run or when config changes
+if "chat_graph" not in st.session_state or st.sidebar.button("Reset Conversation"):
+    if "GROQ_API_KEY" in os.environ:
+        with st.spinner("Initializing chatbot..."):
+            st.session_state.chat_graph = create_chat_graph(system_prompt, selected_model)
+            st.session_state.messages = []  # Clear messages on reset
+            st.success("Chatbot initialized!")
+    else:
+        st.sidebar.error("API key required to initialize chatbot.")
+
+# Display chat history
+for message in st.session_state.messages:
+    if isinstance(message, dict):  # Handle dict format
+        role = message.get("role", "")
+        content = message.get("content", "")
+    else:  # Handle direct string format
+        role = "user" if message.startswith("User: ") else "assistant"
+        content = message.replace("User: ", "").replace("Assistant: ", "")
+    
+    with st.chat_message(role):
+        st.write(content)
+
+# Input for new message
+if "chat_graph" in st.session_state and "GROQ_API_KEY" in os.environ:
+    user_input = st.chat_input("Type your message here...")
+    
+    if user_input:
+        # Display user message
+        with st.chat_message("user"):
+            st.write(user_input)
+        
+        # Add to history
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        
+        # Get response from the chatbot
+        with st.spinner("Thinking..."):
+            # Call the graph with the user's message
+            config = {"configurable": {"thread_id": st.session_state.session_id}}
+            user_message = [HumanMessage(content=user_input)]
+            result = st.session_state.chat_graph.invoke({"messages": user_message}, config)
+            
+            # Extract response
+            response = result["messages"][-1].content
+        
+        # Display assistant response
+        with st.chat_message("assistant"):
+            st.write(response)
+        
+        # Add to history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
+# Add some additional info in the sidebar
+st.sidebar.markdown("---")
+st.sidebar.subheader("About")
+st.sidebar.info(
+    """
+    This chatbot uses LangGraph for maintaining conversation context and 
+    ChatGroq's language models for generating responses. Each conversation 
+    has a unique session ID to maintain history.
+    """
+)
+
+# Download chat history
+if st.sidebar.button("Download Chat History"):
+    import json
+    from datetime import datetime
+    
+    # Convert chat history to downloadable format
+    chat_export = "\n".join([f"{m['role']}: {m['content']}" for m in st.session_state.messages])
+    
+    # Create download button
+    st.sidebar.download_button(
+        label="Download as Text",
+        data=chat_export,
+        file_name=f"chat_history_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt",
+        mime="text/plain"
+    )