components/chat_loop.py

import streamlit as st
from components.generate_chat_stream import generate_chat_stream
from components.stream_handler import stream_handler
from components.show_source import show_source


def chat_loop(session_state, config):
    if prompt := st.chat_input("Search the web..."):
        st.chat_message("user").markdown(prompt)
        session_state.messages.append({"role": "user", "content": prompt})

        chat_stream, links = generate_chat_stream(session_state, prompt, config)

        with st.chat_message("assistant"):
            placeholder = st.empty()
            full_response = stream_handler(
                session_state, chat_stream, prompt, placeholder
            )
            if session_state.rag_enabled:
                show_source(links)

        session_state.history.append([prompt, full_response])
        session_state.messages.append({"role": "assistant", "content": full_response})