ai-lawyer-agent/app/app.py
2025-12-15 23:21:10 +01:00

93 lines
2.9 KiB
Python

from datetime import datetime
import streamlit as st
from core.initialize_agent import assistant_agent
from core.stream_response import SQLiteSession
from app.components.sidebar import add_sidebar
style_chat_message = """
<style>
.st-emotion-cache-1rs7fk9 {
background-color: #D6EDFF;
border-radius: 1rem;
}
.st-emotion-cache-1q1vt2q {
border-radius: 1rem;
}
</style>
"""
def get_time() -> str:
return datetime.now().strftime("%d.%m.%Y %H:%M:%S")
def init_session_state() -> None:
if "messages" not in st.session_state:
st.session_state.messages = []
if "chat_session" not in st.session_state:
st.session_state.chat_session = SQLiteSession(":memory:")
if "show_about" not in st.session_state:
st.session_state.show_about = True
st.markdown(style_chat_message, unsafe_allow_html=True)
def create_app() -> None:
st.set_page_config(
page_title="LawGPT",
page_icon="app/assets/images/title.png",
initial_sidebar_state="collapsed",
layout="centered",
menu_items={
'Get help': None,
'Report a bug': None,
'About': """
This is a cool educational project exploring the creation of an AI agent powered by API keys.
You can learn how to build, interact with, and experiment with AI using real API integration.
"""
}
)
add_sidebar()
init_session_state()
for message in st.session_state.messages:
with st.chat_message(message["role"], avatar=message["avatar"]):
st.markdown(message["content"])
if "time" in message:
st.caption(message["time"])
user_avatar = "app/assets/images/user.png"
assistant_avatar = "app/assets/images/assistant.png"
if request := st.chat_input("Ask anything"):
user_time = get_time()
with st.chat_message(name="user", avatar=user_avatar):
st.markdown(f"{request}")
st.caption(user_time)
user_message = {"role": "user",
"avatar": user_avatar,
"content": request,
"time": user_time}
st.session_state.messages.append(user_message)
with st.chat_message(name="assistant", avatar=assistant_avatar):
with st.spinner("Thinking..."):
try:
response = st.write_stream(assistant_agent(request, st.session_state.chat_session))
except Exception as e:
response = f"️⚠️🌐 Error: {e}"
finally:
assistant_time = get_time()
st.caption(assistant_time)
assistant_message = {"role": "assistant",
"avatar": assistant_avatar,
"content": response,
"time": assistant_time}
st.session_state.messages.append(assistant_message)
st.rerun()