Ameer1606/PostgresPro-Support

1import streamlit as st2import re3import requests4import psycopg25 6from langchain_core.messages import HumanMessage, AIMessage7from src.agent import SupportAgent8from src.config import (9    GROQ_API_KEY,10    QDRANT_URL,11    QDRANT_API_KEY,12    PG_HOST,13    PG_PORT,14    PG_USER,15    PG_PASSWORD,16    PG_DATABASE,17    PG_SSL,18    COMPANY_NAME,19    COMPANY_PHONE,20    COMPANY_EMAIL,21)22 23# -----------------------------24# CONNECTION CHECKS25# -----------------------------26 27def check_qdrant():28    try:29        if not QDRANT_URL:30            return False31 32        headers = {}33        if QDRANT_API_KEY:34            headers["api-key"] = QDRANT_API_KEY35 36        r = requests.get(f"{QDRANT_URL}/collections", headers=headers, timeout=3)37        return r.status_code == 20038    except:39        return False40 41 42def check_postgres():43    try:44        conn = psycopg2.connect(45            host=PG_HOST,46            port=PG_PORT,47            user=PG_USER,48            password=PG_PASSWORD,49            dbname=PG_DATABASE,50            sslmode=PG_SSL,51        )52        conn.close()53        return True54    except:55        return False56 57 58# -----------------------------59# STREAMLIT APP60# -----------------------------61 62def main():63    st.set_page_config(64        page_title=f"{COMPANY_NAME} Agent",65        page_icon="🐘",66        layout="wide",67    )68 69    st.title(f"🐘 {COMPANY_NAME} Support Agent")70    st.markdown(f"**Contact us:** {COMPANY_PHONE} | {COMPANY_EMAIL}")71    st.markdown("AI-powered assistant for PostgreSQL internals and support.")72 73    # -------------------------74    # SIDEBAR STATUS75    # -------------------------76    with st.sidebar:77        st.header("System Status")78 79        qdrant_status = "🟢 Online" if check_qdrant() else "🔴 Offline"80        postgres_status = "🟢 Online" if check_postgres() else "🔴 Offline"81 82        st.markdown(f"**Qdrant (Cloud Vector DB):** {qdrant_status}")83        st.markdown(f"**PostgreSQL (Neon):** {postgres_status}")84 85        if st.button("Reset Chat"):86            st.session_state.messages = []87 88    # -------------------------89    # VALIDATION90    # -------------------------91    if not GROQ_API_KEY:92        st.error("Missing GROQ_API_KEY. Please set it in your .env file.")93        st.stop()94 95    # -------------------------96    # INIT AGENT97    # -------------------------98    if "agent" not in st.session_state:99        st.session_state.agent = SupportAgent()100 101    if "messages" not in st.session_state:102        st.session_state.messages = []103 104    # -------------------------105    # EMPTY STATE106    # -------------------------107    if not st.session_state.messages:108        st.info("Ask anything about PostgreSQL internals, performance, indexing, or troubleshooting.")109 110    # -------------------------111    # CHAT HISTORY112    # -------------------------113    for msg in st.session_state.messages:114        role = "user" if isinstance(msg, HumanMessage) else "assistant"115 116        with st.chat_message(role):117            st.markdown(msg.content)118 119            if role == "assistant":120                citations = set(121                    re.findall(r"Source:\s*([^,]+),\s*Page:\s*(\d+)", msg.content)122                )123 124                if citations:125                    with st.expander("View Document References"):126                        for source, page in citations:127                            st.write(f"- **{source}** (Page {page})")128 129    # -------------------------130    # CHAT INPUT131    # -------------------------132    if prompt := st.chat_input("How can I help you with PostgreSQL?"):133 134        st.session_state.messages.append(HumanMessage(content=prompt))135 136        with st.chat_message("user"):137            st.markdown(prompt)138 139        with st.chat_message("assistant"):140            with st.spinner("Thinking..."):141                try:142                    response = st.session_state.agent.run(143                        prompt,144                        st.session_state.messages,145                    )146 147                    st.markdown(response)148 149                    citations = set(150                        re.findall(r"Source:\s*([^,]+),\s*Page:\s*(\d+)", response)151                    )152 153                    if citations:154                        with st.expander("View Document References"):155                            for source, page in citations:156                                st.write(f"- **{source}** (Page {page})")157 158                    st.session_state.messages.append(AIMessage(content=response))159 160                except Exception as e:161                    st.error(f"Error: {str(e)}")162 163 164if __name__ == "__main__":165    main()