PABPAT/TCI_Shield

1import streamlit as st2import re3import os4import tempfile5from tci_agent import reset_session, SYSTEM_PROMPT, nova_lite6from tci_agent import get_progress, set_buyer_count, collect_business_info7from tci_agent import collect_buyer_info, upload_financial_document8from tci_agent import collect_financial_data, run_underwriting9from tci_agent import generate_policy_options, issue_policy10from strands import Agent11 12 13def clean(text: str) -> str:14    return re.sub(r"<thinking>.*?</thinking>", "", str(text), flags=re.DOTALL).strip()15 16 17st.title("Text Chat")18st.write("Chat with Alex, your AI underwriter.")19 20if "chat_history" not in st.session_state:21    st.session_state.chat_history = []22 23if "agent" not in st.session_state:24    reset_session()25    st.session_state.agent = Agent(26        model=nova_lite,27        system_prompt=SYSTEM_PROMPT,28        callback_handler=None,29        tools=[30            get_progress,31            set_buyer_count,32            collect_business_info,33            collect_buyer_info,34            upload_financial_document,35            collect_financial_data,36            run_underwriting,37            generate_policy_options,38            issue_policy,39        ],40    )41    opening = clean(st.session_state.agent(42        "Hello, I would like to apply for trade credit insurance."43    ))44    st.session_state.chat_history.append({45        "role": "assistant",46        "content": opening47    })48 49for msg in st.session_state.chat_history:50    with st.chat_message(msg["role"]):51        st.write(msg["content"])52 53user_input = st.chat_input("Type your message...")54 55if user_input:56    st.session_state.chat_history.append({57        "role": "user",58        "content": user_input59    })60    with st.chat_message("user"):61        st.write(user_input)62 63    response_text = clean(st.session_state.agent(user_input))64 65    st.session_state.chat_history.append({66        "role": "assistant",67        "content": response_text68    })69    with st.chat_message("assistant"):70        st.write(response_text)71 72st.divider()73st.subheader("Upload Financial Statement")74uploaded_file = st.file_uploader(75    "Upload your financial statement (PDF, DOCX, XLSX, CSV)",76    type=["pdf", "docx", "xlsx", "csv", "txt"]77)78 79if uploaded_file:80    with tempfile.NamedTemporaryFile(81        delete=False,82        suffix=os.path.splitext(uploaded_file.name)[1]83    ) as tmp:84        tmp.write(uploaded_file.read())85        tmp_path = tmp.name86 87    if st.button("Extract Financial Data from Document"):88        with st.spinner("Reading document with Nova Multimodal..."):  # type: ignore89            response_text = clean(st.session_state.agent(90                f"The customer has uploaded a financial statement. "91                f"Please extract the financial figures using the upload_financial_document tool. "92                f"File path: {tmp_path}"93            ))94            st.success("Financial data extracted!")95            st.session_state.chat_history.append({96                "role": "assistant",97                "content": response_text98            })99            st.rerun()100 101if st.button("Start New Conversation"):102    reset_session()103    st.session_state.chat_history = []104    del st.session_state.agent105    st.rerun()