import streamlit as st
from middlewares.utils import gen_augmented_prompt_via_websearch
from middlewares.chat_client import chat


def generate_chat_stream(session_state, prompt, config):
    # 1. augments prompt according to the template
    # 2. returns chat_stream and source links
    # 3. chat_stream and source links are used by stream_handler and show_source
    chat_bot_dict = config["CHAT_BOTS"]
    links = []
    if session_state.rag_enabled:
        with st.spinner("Fetching relevent documents from Web...."):
            prompt, links = gen_augmented_prompt_via_websearch(
                prompt=prompt,
                pre_context=session_state.pre_context,
                post_context=session_state.post_context,
                pre_prompt=session_state.pre_prompt,
                post_prompt=session_state.post_prompt,
                vendor=session_state.search_vendor,
                top_k=session_state.top_k,
                n_crawl=session_state.n_crawl,
                pass_prev=session_state.pass_prev,
                prev_output=session_state.history[-1][1],
            )

    with st.spinner("Generating response..."):
        chat_stream = chat(
            prompt,
            session_state.history,
            chat_client=chat_bot_dict[session_state.chat_bot],
            temperature=session_state.temp,
            max_new_tokens=session_state.max_tokens,
        )

    return chat_stream, links