app.py

# app.py
import os
import streamlit as st
from ui import ChatUI
from client import Qwen3Client

from server.web_configs import WEB_CONFIGS, API_CONFIG
from utils.api import get_rag_api


class SessionState:
    def init_state(self):
        """Initialize session state variables."""
        st.session_state['assistant'] = []
        st.session_state['user'] = []

        # action_list = [
        #     GoogleSearch(),
        #     WeatherQuery(),
        #     ItineraryPlan(),
        # ]
        # st.session_state['plugin_map'] = {
        #     action.name: action
        #     for action in action_list
        # }
        st.session_state['model_map'] = {}
        st.session_state['model_selected'] = None
        st.session_state['plugin_actions'] = set()
        st.session_state['history'] = []

    def clear_state(self):
        """Clear the existing session state."""
        st.session_state['assistant'] = []
        st.session_state['user'] = []
        st.session_state['model_selected'] = None
        st.session_state['file'] = set()
        if 'chatbot' in st.session_state:
            st.session_state['chatbot']._session_history = []


def main():
    # Initialize Streamlit UI and setup sidebar
    if 'ui' not in st.session_state:
        session_state = SessionState()
        session_state.init_state()
        st.session_state['ui'] = ChatUI(session_state)
    else:
        st.set_page_config(
            layout='wide',
            page_title='RAG',
            page_icon='./docs/imgs/avic.png')
        st.header('RAG聊天机器人', divider='rainbow')
    
    ui = st.session_state["ui"]
    cfg, uploaded_file = ui.sidebar()

    # if 'chatbot' not in st.session_state or model != st.session_state[
    #         'chatbot']._llm:
    st.session_state['chatbot'] = st.session_state['ui'].initialize_chatbot(cfg, "llm")
    # st.session_state['chatbot'] = st.session_state['ui'].initialize_chatbot(cfg, "agent")
    st.session_state['session_history'] = []
    ui.render_history()
    for prompt, agent_return in zip(st.session_state['user'],
                                    st.session_state['assistant']):
        st.session_state['ui'].render_user(prompt)
        st.session_state['ui'].render_assistant_final(agent_return)

    user_input = ui.get_user_input()
    if not user_input:
        return
    with st.container():
        st.session_state['ui'].render_user(user_input)
    st.session_state['user'].append(user_input)
    # Add file uploader to sidebar
    if (uploaded_file and uploaded_file.name not in st.session_state['file']):
        st.session_state['file'].add(uploaded_file.name)
        file_bytes = uploaded_file.read()
        file_type = uploaded_file.type
        if 'image' in file_type:
            st.image(file_bytes, caption='Uploaded Image')
        elif 'video' in file_type:
            st.video(file_bytes, caption='Uploaded Video')
        elif 'audio' in file_type:
            st.audio(file_bytes, caption='Uploaded Audio')
        # Save the file to a temporary location and get the path
        postfix = uploaded_file.name.split('.')[-1]
        # prefix = str(uuid.uuid4())
        prefix = hashlib.md5(file_bytes).hexdigest()
        filename = f'{prefix}.{postfix}'
        file_path = os.path.join(root_dir, filename)
        with open(file_path, 'wb') as tmpfile:
            tmpfile.write(file_bytes)
        file_size = os.stat(file_path).st_size / 1024 / 1024
        file_size = f'{round(file_size, 2)} MB'
        st.write(f'File saved at: {file_path}')
        user_input = [
            dict(role='user', content=user_input),
            dict(
                role='user',
                content=json.dumps(dict(path=file_path, size=file_size)),
                name='file')
        ]
    if isinstance(user_input, str):
        user_input = [dict(role='user', content=user_input)]
    # ui.render_user([dict(role='user', content=str(user_input))])
    # 先采用RAG查看是否能检索到相似文本
    user_input = get_rag_api(user_input)
    st.session_state.messages.append(user_input[-1])
    # client = Qwen3Client(cfg["server_url"])

    client = st.session_state['chatbot']
    if type(client) == Qwen3Client:
        if cfg["stream"]:
            for chunk in client.chat_completions(
                messages=st.session_state.messages,
                model=cfg["model"],
                temperature=cfg["temperature"],
                top_p=cfg["top_p"],
                top_k=cfg["top_k"],
                max_tokens=cfg["max_tokens"],
                presence_penalty=cfg["presence_penalty"],
                frequency_penalty=cfg["frequency_penalty"],
                stream=True,
                enable_thinking=cfg["enable_thinking"],
            ):
                delta = (
                    chunk.get("choices", [{}])[0]
                    .get("delta", {})
                    .get("content", "")
                )
                ui.render_assistant_stream(delta)

            st.session_state.messages.append(
                {
                    "role": "assistant",
                    "content": st.session_state.assistant_text
                }
            )
            del st.session_state.assistant_text
            del st.session_state.assistant_placeholder

        else:
            result = client.chat_completions(
                messages=st.session_state.messages,
                model=cfg["model"],
                temperature=cfg["temperature"],
                top_p=cfg["top_p"],
                top_k=cfg["top_k"],
                max_tokens=cfg["max_tokens"],
                presence_penalty=cfg["presence_penalty"],
                frequency_penalty=cfg["frequency_penalty"],
                stream=False,
                enable_thinking=cfg["enable_thinking"],
            )
            content = result["choices"][0]["message"]["content"]
            ui.render_assistant_final(content)
            st.session_state.messages.append(
                {"role": "assistant", "content": content}
            )
    else:
        for event in client.stream({"messages": user_input}):
            # if "files" in event:
            #     print ("*"*80)
            #     print (event["files"])
            # event["messages"][-1].pretty_print()
            print(event)


if __name__ == "__main__":
    root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    root_dir = os.path.join(root_dir, 'tmp_dir')
    os.makedirs(root_dir, exist_ok=True)
    main()