| import streamlit as st |
| from dataclasses import dataclass |
| from langchain.schema import HumanMessage |
| from agent import chat_agent |
|
|
|
|
| @dataclass |
| class Message: |
| actor: str |
| payload: str |
|
|
| USER = "user" |
| ASSISTANT = "ai" |
| MESSAGES = "messages" |
|
|
| |
| st.set_page_config(page_title="AI Chatbot", page_icon="🤖") |
| st.title("🤖 First AI Chatbot") |
|
|
| |
| |
| if MESSAGES not in st.session_state: |
| st.session_state[MESSAGES] = [Message(actor=ASSISTANT, payload="Hi! How can I help you?")] |
|
|
| msg: Message |
| for msg in st.session_state[MESSAGES]: |
| st.chat_message(msg.actor).write(msg.payload) |
|
|
|
|
| |
| if prompt := st.chat_input("Type your message..."): |
| st.session_state[MESSAGES].append(Message(actor=USER, payload=prompt)) |
| st.chat_message(USER).write(prompt) |
| result = chat_agent.invoke({ |
| "messages": [HumanMessage(content=prompt)] |
| }) |
| response: str = result["messages"][-1].content |
| st.session_state[MESSAGES].append(Message(actor= ASSISTANT, payload= response)) |
| st.chat_message(ASSISTANT).write(response) |
