Spaces:
Sleeping
Sleeping
File size: 2,559 Bytes
cce43bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
"""
Session management for the AI Database Assistant.
"""
import streamlit as st
from typing import List, Dict, Any
class SessionManager:
"""Manages session state and message handling."""
@staticmethod
def initialize_session() -> None:
"""Initialize session state variables."""
if "messages" not in st.session_state:
st.session_state["messages"] = []
@staticmethod
def get_messages() -> List[Dict[str, Any]]:
"""Get all messages from session state."""
return st.session_state.get("messages", [])
@staticmethod
def add_user_message(content: str) -> None:
"""Add a user message to the session."""
st.session_state["messages"].append({
"role": "user",
"content": content
})
@staticmethod
def add_thinking_message() -> None:
"""Add a thinking placeholder message."""
st.session_state["messages"].append({
"role": "assistant",
"content": "🤔 Thinking...",
"is_placeholder": True
})
@staticmethod
def replace_last_message(content: str, data: List[Dict] = None,
chart: str = None, is_error: bool = False) -> None:
"""Replace the last message with final response."""
st.session_state["messages"][-1] = {
"role": "assistant",
"content": content,
"data": data or [],
"chart": chart,
"is_error": is_error,
}
@staticmethod
def is_ai_thinking() -> bool:
"""Check if AI is currently thinking (has placeholder message)."""
messages = SessionManager.get_messages()
if not messages:
return False
last_message = messages[-1]
return (
last_message["role"] == "assistant" and
last_message.get("is_placeholder", False)
)
@staticmethod
def get_last_user_message() -> str:
"""Get the content of the last user message."""
messages = SessionManager.get_messages()
if len(messages) >= 2:
return messages[-2]["content"]
return ""
@staticmethod
def has_pending_response() -> bool:
"""Check if there's a pending response to process."""
messages = SessionManager.get_messages()
return (
messages and
messages[-1].get("is_placeholder") and
len(messages) >= 2 and
messages[-2]["role"] == "user"
)
|