import os import glob import uuid import asyncio # import trace_config import logging import streamlit as st from aagents.orchestrator_agent import orchestrator_agent from agents import Runner, trace, SQLiteSession from agents.exceptions import InputGuardrailTripwireTriggered # from langsmith import traceable from traceloop.sdk import Traceloop from opentelemetry.sdk.trace import Span # --- Monkeypatch to fix "Invalid type Omit" errors --- # This filters out 'NotGiven'/'Omit' values from OpenAI that crash the OTel exporter _original_set_attribute = Span.set_attribute def _safe_set_attribute(self, key, value): # Check string representation of type to avoid importing specific internal types type_str = str(type(value)) if "Omit" in type_str or "NotGiven" in type_str: return self return _original_set_attribute(self, key, value) Span.set_attribute = _safe_set_attribute # ----------------------------------------------------- Traceloop.init( disable_batch=True, api_key="tl_1c19b8e8fcfd411fb9fcdb02d381faef" ) # ----------------------------- # Configuration & Utils # ----------------------------- st.set_page_config( page_title="AI Assistant", layout="wide", page_icon="🤖" ) def load_prompts(folder="prompts"): prompts = [] prompt_labels = [] if os.path.exists(folder): for file_path in glob.glob(os.path.join(folder, "*.txt")): with open(file_path, "r", encoding="utf-8") as f: content = f.read().strip() if content: prompts.append(content) prompt_labels.append(os.path.basename(file_path).replace("_", " ").replace(".txt", "").title()) return prompts, prompt_labels prompts, prompt_labels = load_prompts() # ----------------------------- # Session State # ----------------------------- if "messages" not in st.session_state: st.session_state.messages = [] if "ai_session_id" not in st.session_state: st.session_state.ai_session_id = str(uuid.uuid4()) # Persistent SQLite session if "ai_session" not in st.session_state: st.session_state.ai_session = SQLiteSession(f"conversation_{st.session_state.ai_session_id}.db") session = st.session_state.ai_session # ----------------------------- # Premium Styling # ----------------------------- st.markdown(""" """, unsafe_allow_html=True) # ----------------------------- # Logic # ----------------------------- # @traceable(name="chatbot") async def get_ai_response(prompt: str) -> str: try: agent = orchestrator_agent # Ensure session is valid current_session = st.session_state.ai_session current_session = st.session_state.ai_session with trace("Chatbot Agent Run"): # Keep existing custom trace wrapper # Run agent result = await Runner.run(agent, prompt, session=current_session) return result.final_output except InputGuardrailTripwireTriggered as e: reasoning = getattr(e, "reasoning", None) \ or getattr(getattr(e, "output", None), "reasoning", None) \ or getattr(getattr(e, "guardrail_output", None), "reasoning", None) \ or "Guardrail triggered, but no reasoning provided." return f"⚠️ **Guardrail Blocked Input**\n\n{reasoning}" except Exception as e: return f"❌ **Error**: {str(e)}" # ----------------------------- # Sidebar - Quick Actions # ----------------------------- with st.sidebar: st.markdown("### ⚡ Quick Starters") st.markdown("Select a prompt to start:") # We use a trick with st.button to act as input triggers # If a button is clicked, we'll handle it in the main loop logic selected_prompt = None for idx, prompt_text in enumerate(prompts): label = prompt_labels[idx] if idx < len(prompt_labels) else f"Prompt {idx+1}" if st.button(label, key=f"sidebar_btn_{idx}", use_container_width=True): # Reset conversation st.session_state.messages = [] st.session_state.ai_session_id = str(uuid.uuid4()) # Recreate session object with new ID st.session_state.ai_session = SQLiteSession(f"conversation_{st.session_state.ai_session_id}.db") selected_prompt = prompt_text st.markdown("---") if st.button("🗑️ Clear Conversation", use_container_width=True): st.session_state.messages = [] st.rerun() # ----------------------------- # Main Content # ----------------------------- # Hero Banner (Always visible & Sticky) st.markdown(""" """, unsafe_allow_html=True) # Display Chat History for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"], unsafe_allow_html=True) # Chat Input Handling # We handle both the chat input widget and the sidebar selection here if prompt := (st.chat_input("Type your message...") or selected_prompt): # User Message st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): st.markdown(prompt) # Assistant Response with st.chat_message("assistant"): with st.spinner("Thinking..."): response_text = asyncio.run(get_ai_response(prompt)) st.markdown(response_text, unsafe_allow_html=True) st.session_state.messages.append({"role": "assistant", "content": response_text}) # If it was a sidebar click, we need to rerun to clear the selection state potentially, # but st.chat_input usually handles focus. With buttons, a rerun happens automatically # but we want to make sure the input box is cleared (which 'selected_prompt' doesn't use). if selected_prompt: st.rerun()