import json import sys from pathlib import Path if sys.version_info >= (3, 11): import tomllib else: import tomli as tomllib import requests import snowflake.connector import streamlit as st _config_path = Path(__file__).parent / "config.toml" with open(_config_path, "rb") as _f: _cfg = tomllib.load(_f) AGENT_DATABASE = _cfg["agent"]["database"] AGENT_SCHEMA = _cfg["agent"]["schema"] AGENT_NAME = _cfg["agent"]["name"] st.set_page_config(page_title=_cfg["ui"]["title"], page_icon=":musical_note:", layout="wide") st.title(_cfg["ui"]["title"]) st.caption(_cfg["ui"]["caption"]) @st.cache_resource def get_snowflake_connection(): return snowflake.connector.connect( account=st.secrets["connections"]["snowflake"]["account"], user=st.secrets["connections"]["snowflake"]["user"], authenticator=st.secrets["connections"]["snowflake"]["authenticator"], role=st.secrets["connections"]["snowflake"].get("role"), database=st.secrets["connections"]["snowflake"].get("database"), schema=st.secrets["connections"]["snowflake"].get("schema"), ) def get_token_and_host(): conn = get_snowflake_connection() return conn.rest.token, conn.host def create_thread(token: str, host: str) -> str: url = f"https://{host}/api/v2/cortex/threads" headers = { "Authorization": f'Snowflake Token="{token}"', "Content-Type": "application/json", } resp = requests.post(url, headers=headers, json={}, verify=False) resp.raise_for_status() return str(resp.json().get("thread_id", "")) def get_agent_url(host: str) -> str: return f"https://{host}/api/v2/databases/{AGENT_DATABASE}/schemas/{AGENT_SCHEMA}/agents/{AGENT_NAME}:run" def call_agent(question: str, thread_id: str, parent_message_id: int): """Call the agent via SSE. Returns (text, charts, assistant_message_id).""" token, host = get_token_and_host() url = get_agent_url(host) headers = { "Authorization": f'Snowflake Token="{token}"', "Content-Type": "application/json", } payload = { "messages": [ { "role": "user", "content": [{"type": "text", "text": question}], } ], "thread_id": thread_id, "parent_message_id": parent_message_id, } resp = requests.post(url, headers=headers, json=payload, stream=True, verify=False) if resp.status_code != 200: error_text = f"Error: {resp.status_code} — {resp.text}" return error_text, [], 0 text_parts = [] charts = [] assistant_message_id = 0 event_type = None for line in resp.iter_lines(): if not line: continue decoded = line.decode("utf-8") if decoded.startswith("event: "): event_type = decoded[7:].strip() elif decoded.startswith("data: "): if event_type == "done": break try: data = json.loads(decoded[6:]) except json.JSONDecodeError: continue if event_type == "metadata": meta = data.get("metadata", {}) if meta.get("role") == "assistant": mid = meta.get("message_id") if mid: assistant_message_id = mid elif event_type == "response.text.delta": text_chunk = data.get("text", "") if text_chunk: text_parts.append(text_chunk) yield {"type": "text", "text": text_chunk} elif event_type == "response.chart": spec_str = data.get("chart_spec") if spec_str: try: spec = json.loads(spec_str) if isinstance(spec_str, str) else spec_str charts.append(spec) yield {"type": "chart", "spec": spec} except json.JSONDecodeError: pass elif event_type == "response": if not assistant_message_id: mid = data.get("metadata", {}).get("assistant_message_id") if mid: assistant_message_id = mid full_text = "".join(text_parts) yield {"type": "done", "text": full_text, "charts": charts, "assistant_message_id": assistant_message_id} # --- Session state initialization --- if "messages" not in st.session_state: st.session_state.messages = [] if "thread_id" not in st.session_state: st.session_state.thread_id = None if "parent_message_id" not in st.session_state: st.session_state.parent_message_id = 0 def ensure_thread(): if not st.session_state.thread_id: token, host = get_token_and_host() st.session_state.thread_id = create_thread(token, host) st.session_state.parent_message_id = 0 # --- Sidebar --- with st.sidebar: if st.button("New Conversation"): st.session_state.messages = [] st.session_state.thread_id = None st.session_state.parent_message_id = 0 st.rerun() # --- Suggestion chips --- SUGGESTIONS = {s["label"]: s["prompt"] for s in _cfg["ui"]["suggestions"]} if not st.session_state.messages: selected = st.pills("Try asking:", list(SUGGESTIONS.keys()), label_visibility="collapsed") if selected: st.session_state.messages.append({"role": "user", "content": SUGGESTIONS[selected]}) st.rerun() # --- Render chat history --- for msg in st.session_state.messages: with st.chat_message(msg["role"]): st.markdown(msg["content"]) for chart_spec in msg.get("charts", []): st.vega_lite_chart(chart_spec, use_container_width=True) # --- Handle new input --- if prompt := st.chat_input(_cfg["ui"]["chat_placeholder"]): ensure_thread() st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): st.markdown(prompt) with st.chat_message("assistant"): text_placeholder = st.empty() streamed_text = "" collected_charts = [] for event in call_agent(prompt, st.session_state.thread_id, st.session_state.parent_message_id): if event["type"] == "text": streamed_text += event["text"] text_placeholder.markdown(streamed_text + "▌") elif event["type"] == "chart": collected_charts.append(event["spec"]) elif event["type"] == "done": text_placeholder.markdown(streamed_text) for chart_spec in collected_charts: st.vega_lite_chart(chart_spec, use_container_width=True) st.session_state.parent_message_id = event["assistant_message_id"] st.session_state.messages.append({ "role": "assistant", "content": streamed_text, "charts": collected_charts, })