app.py

"""
1. 用Mermaid Mardown来画各种图，包括：脑图、流程图、客户历程图、关系图等。
1. 目前ChatGPT级别和百度文星一言等级别的大模型都可以比较好的输出Mermaid Markdown格式图。
1. 注意Mermaid代码中，不需要任何的’‘’， ````，或者""，这里只需要直接写Mermaid代码即可。
1. 完成了regenerate功能。
1. 完成了历史记忆功能，可以对图进行持续优化。但是不能保存中间结果。

1. 已知问题和解决方法：
    1. 已解决。尝试了多个方法，目前无法把HTML展现的内容用图片保存，目前不行。直接保存MHTML文件。
        1. HTML文件本身可以保存。
        1. 用html2image库，可以把HTML文件转换为图片，但是无法在streamlit中展示（内容为空白）。
        1. 可以想到的workaround：把HTML页面保存下来。
    1. HTML的高度只能手工制定，无法自动调整。可能垂直的页面会被截断。已经解决，通过设定scrolling=True, 可以实现滚动条。高度不在重要。
    1. 已解决。点击下载button后，页面会刷新。采用一个特殊的extension，https://github.com/PaleNeutron/streamlit-ext
"""
####TODO：

import json
import pprint
import streamlit as st
import chatgpt

import markdown
import md_mermaid
# from streamlit import components
import requests
import re
from openai import OpenAI
import streamlit_authenticator as stauth


import streamlit_ext as ste ##TODO： 为了点击download button后保持页面。

from datetime import datetime
from pytz import timezone

### streamlit app title
st.set_page_config(layout="wide", page_icon='llm_icon.png') ## 必须是第一行
st.title(f"大语言模型 - 体系图 | 框架图 | 逻辑图 | 流程图 - 辅助设计中心", anchor='Title')
st.subheader("AI Flowchart - Mindmap - Relation Diagram Design for Professionals")
st.info('如果输出图例时遇见任何问题（如：syntax error）或者不满意当前结果，请在左侧重新提交您的问题即可。一般建议至少尝试3-10次。')


## toast effect. 
# msg = st.toast('程序正在启动中，请稍等...',)
# msg.toast('大语言模型成功加载!', icon = "🥞")
# st.toast('大语言模型成功加载!', icon = "🥞")
# st.divider()
# st.markdown('_说明：如果输出图例遇见有任何问题，请刷新页面再重新提交您的问题即可！_')
# st.snow() ##可以在页面上显示雪花效果。
# st.balloons() ##可以在页面上显示气球效果。


### system prompt设置
openai_client = OpenAI()

# user_input = ""
system_prompt = f"""你是一个Mermaid Markdown方面的设计专家。你需要根据'我的要求'用Mermaid Markdown来画出对应的图。你仅需要提供Mermaid Markdown格式的代码，你不允许输出任何说明、解释或者提示的内容。"""
# system_prompt = f"""你是一个Mermaid Markdown方面的设计专家。你需要根据'我的要求'用Mermaid Markdown来画出对应的图。你仅需要提供Mermaid Markdown格式的代码，你不允许输出任何说明、解释或者提示的内容。

# '我的要求'如下：{user_input}"""

### 用户输入框
# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = [] ### original code here.
    # st.session_state.messages = [{"role": "system", "content": "你是一个Mermaid Markdown专家"}]
    # st.session_state.messages = [{"role": "system", "content": system_prompt}]

# st.session_state.messages = [{"role": "system", "content": '你是一个Mermaid Markdown方面的设计专家。'},{"role": "user", "content": ''}]


user_input = st.chat_input("说点什么吧...") ### original code here.
# if user_input:
#     with st.chat_message("user"):
#         st.markdown(user_input) ## 这里只需要输出用户的prompt，而不是total prompt。
#         st.session_state.messages.append({"role": "user", "content": user_input})

print('st.session_state.messages:')
pprint.pprint(st.session_state.messages)
# print('user_input:', json.dumps(user_input))
print('user_input now:')
pprint.pprint(user_input)
# print('user_input now', user_input)
# st.session_state.messages.append({"role": "user", "content": user_input})
### 设定一个历史信息的列表，用于存储用户的输入。为了regenerate按钮的功能。
# hist_msg = []



### 配置前端有关函数
# user_input = None
def clear_all():
    st.session_state.conversation = None
    st.session_state.chat_history = None
    st.session_state.messages = []
    message_placeholder = st.empty()
    st.session_state["my_question"] = None
    return None



### 重新生成按钮, regenerate, rerun（streamlit中的rerun是把整个页面重新加载一次？）, resubmit. 
def regenerate():
    html_file = ""
    ## 因为可能没有历史，第一次的时候，所以需要处理异常。
    try:
        print('st.session_state.messages inside REGENERATE function:')
        pprint.pprint(st.session_state.messages)
        st.session_state.messages = st.session_state.messages[:-1]
        html_file = main(input=st.session_state.messages[0]['content']) ### original code here.
        # html_file = main(input=st.session_state.messages[-1]['content']) ### original code here.
    except Exception as e:
        print('Error:', e)
        pass
        
    return html_file

### authentication with a local yaml file.
import yaml
from yaml.loader import SafeLoader
with open('./config.yaml') as file:
    config = yaml.load(file, Loader=SafeLoader)
authenticator = stauth.Authenticate(
    config['credentials'],
    config['cookie']['name'],
    config['cookie']['key'],
    config['cookie']['expiry_days'],
    config['preauthorized']
)

user, authentication_status, username = authenticator.login('main')
# user, authentication_status, username = authenticator.login('用户登录', 'main')


### streamlit sidebar
if authentication_status:
    with st.sidebar:
        st.markdown(
            """
            <style>
            [data-testid="stSidebar"][aria-expanded="true"]{
                min-width: 450px;
                max-width: 450px;
            }
            """,
            unsafe_allow_html=True,
        )
        ### siderbar的题目。
        ### siderbar的题目。
        # st.header(f'**大语言模型专家系统工作设定区**')
        # st.header(f'**系统控制面板** ')
        st.header(f'**欢迎 **{username}** 来到人工智能的世界** ♠')
        # st.header(f'**欢迎 **{username}** 使用本系统** ') ## 用户登录显示。
        st.write(f'_Large Language Model System Environment_')
        authenticator.logout('登出', 'sidebar')
        # st.divider()
        
        st.sidebar.button("清除记录，重启一轮新对话", on_click=clear_all, use_container_width=True, type='primary')
        re_btn = st.sidebar.button("重新生成答案", use_container_width=True, type='secondary')
        # if re_btn:
        #     regenerate()
        # re_btn = st.sidebar.button("重新生成答案", on_click=regenerate, use_container_width=True, type='primary')

        ## 在sidebar上的三个分页显示，用st.tabs实现。
        # tab_1, tab_2, tab_4 = st.tabs(['使用须知', '模型参数', '角色设定'])
        tab_1, tab_2, tab_3, tab_4 = st.tabs(['基本介绍', '大模型参数', '提示词示例', '使用技巧'])

        # with st.expander(label='**使用须知**', expanded=False):
        with tab_1:
            # st.markdown("#### 快速上手指南")
            # with st.text(body="说明"):
            #     st.markdown("* 重启一轮新对话时，只需要刷新页面（按Ctrl/Command + R）即可。")
            with st.text(body="说明"):
                st.markdown("""* **使用大型语言模型设计体系图、框架图、逻辑图、流程图、关系图主要有以下几个步骤:**

    1. **明确目的和要表达的内容**
    在开始之前,你需要明确你想用图表来表达什么内容, 以及它的目的是什么。这将为接下来的步骤提供指导。

    2. **收集并组织相关信息**
    根据你的目的,收集所有相关的信息、数据和要点。将它们按合理的方式分类和组织, 为后续生成图表做好准备。

    3. **结构化输入**
    将组织好的内容转化为语言模型可以理解的结构化输入。比如使用简单描述、列表、树状结构等形式。

    4. **生成参考版本**
    将结构化输入提交给大型语言模型, 让它尝试生成供参考的流程图、脑图或关系图。

* 需要注意的是, 尽管大型语言模型可以生成很好的草图和初始版本, 但最终结果的质量仍然取决于你对目标和输入的描述质量。人工审查和调整是必不可少的环节。模型只是辅助工具, 无法完全替代人的判断和创造力。""")

        ## 大模型参数
        # with st.expander(label='**大语言模型参数**', expanded=True):
        with tab_2:
            max_tokens = st.slider(label='Max_Token（生成结果时最大字数）', min_value=100, max_value=4096, value=2048, step=100)
            temperature = st.slider(label='Temperature (温度)', min_value=0.0, max_value=1.0, value=0.8, step=0.1)
            top_p = st.slider(label='Top_P (核采样)', min_value=0.0, max_value=1.0, value=0.6, step=0.1)
            frequency_penalty = st.slider(label='Frequency Penalty (重复度惩罚因子)', min_value=-2.0, max_value=2.0, value=1.0, step=0.1)
            presence_penalty = st.slider(label='Presence Penalty (控制主题的重复度)', min_value=-2.0, max_value=2.0, value=1.0, step=0.1)

        ## reset password widget
        # try:
        #     if authenticator.reset_password(st.session_state["username"], 'Reset password'):
        #         st.success('Password modified successfully')
        # except Exception as e:
        #     st.error(e)

        # with st.header(body="欢迎"):
        #     st.markdown("# 欢迎使用大语言模型商业智能中心")
        # with st.expander(label=("**重要的使用注意事项**"), expanded=True):
        # with st.container():

        ##NOTE： 在SQL场景去不需要展示这些提示词。
        with tab_3:
            # st.write("#### Prompt提示词参考资料")
            st.code(body="你输出一个复杂的电子商务流程示意图，必须包含客户投诉退款环节。", language='plaintext')
            st.code(body="给我一个复杂的业务需求分析流程图。", language='plaintext')
            st.code(body="完整的Hermes客户体验流程图。", language='plaintext')
            st.code(body="完整的苹果公司线下实体店客户历程图。", language='plaintext')
            st.code(body="绘制一个横向的详细的IT运维流程示例。", language='plaintext')
            st.code(body="画一个汽车4S店的用户历程的复杂流程（至少要包含：客户邀约，客户试驾等）。", language='plaintext')
            st.code(body="绘制一张保险项目的甘特图。", language='plaintext')
            st.code(body="画一个用户历程的复杂流程。", language='plaintext')
            st.code(body="绘制一张复杂的绩效管理脑图。", language='plaintext')
            st.code(body="给我一个复杂的客户体验脑图示例。", language='plaintext')
            st.code(body="你给我一个全面质量管理的头脑风暴脑图。", language='plaintext')
            st.code(body="给我一个全面且复杂的精益管理流程图，图上需要每一步的说明。", language='plaintext')
            st.code(body="你做一个青少儿教育培训机构的完整电话邀约流程。", language='plaintext')
            st.code(body="在上面流程的基础上，继续细化”系统设计“部分的内容。", language='plaintext')
        
        with tab_4:
            st.markdown('''
    1. 使用`下载图例`按钮，可以直接下载MHTML格式的文件，大部分浏览器支持直接打开。

    1. 使用`下载代码`按钮，可以下载原始的Markdown代码，然后可以在类似 `https://mermaid.live/` 网站手动编辑。

    1. 目前支持持续修改流程版本，但不建议过度修改。修改的中间过程内容不会被保存。
    ''')


elif authentication_status == False:
    st.error('⛔ 用户名或密码错误！')
elif authentication_status == None:
    st.warning('⬆ 请先登录！')


### 得到当前的时间
def get_current_time():
    beijing_tz = timezone('Asia/Shanghai')
    beijing_time = datetime.now(beijing_tz)
    current_time = beijing_time.strftime('%H:%M:%S')
    return current_time

## 显示Mermaid图的核心函数
def mermaid(code: str):
    # st.write(code) ### 检查输入的Mermaid Markdown代码。
    from streamlit import components
    
    ### 以下返回的是HTML文件。
    html_file = components.v1.html(
        f"""
        <pre class="mermaid">
            {code}
        </pre>

        <script type="module">
            import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs';
            
            mermaid.initialize({{ startOnLoad: true }});
        </script>
        """
    ,
    height=400, ##NOTE: 这里可以有效改变HTML页面的高度。
    scrolling=True, ##NOTE: 可以实现滚动条。高度不在重要。
    )
    
    ## 以下是以字符串的形式返回HTML代码。
    html_file = str(f"""
        <pre class="mermaid">
            {code}
        </pre>

        <script type="module">
            import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs';
            
            mermaid.initialize({{ startOnLoad: true }});
        </script>
        """)
    
    ## 以下是以字符串的形式, 且只返回code部分，为了可以在https://mermaid.live/ 编辑。
    mermaid_str = code
    
    return html_file, mermaid_str
    # return html_file


def markdown_chart(input:str):
    # """answer general questions.""" 
    # final_prompt = f"""你是一个Mermaid Markdown方面的设计专家。你需要根据'我的要求'用Mermaid Markdown来画出对应的图。你仅需要提供Mermaid Markdown格式的代码，你不允许输出任何说明、解释或者提示的内容。
    # '我的要求'如下：{user_input}"""
    
    # print('user_input now', user_input)
    # st.session_state.messages.append(
    #                 {"role": "user", "content": user_input})
    # final_prompt = system_prompt + user_input

    if input:
        # st.session_state.messages.append({"role": "user", "content": input})
        # with st.chat_message("user"):
        #     st.markdown(user_input) ## 这里只需要输出用户的prompt，而不是total prompt。

        with st.chat_message("assistant", avatar="./llm_icon.png"):
            message_placeholder = st.empty()
            full_response = ""
            # llm_response = chatgpt.chatgpt(user_prompt=final_prompt) ### original code here.
            print('st.session_state.messages:')
            pprint.pprint(st.session_state.messages)
            # llm_response = chatgpt.chatgpt(user_prompt=st.session_state.messages) ### original code here.
            chatgpt_response = openai_client.chat.completions.create(model="gpt-3.5-turbo-16k",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": json.dumps(st.session_state.messages, ensure_ascii=False)}, ##NOTE：注意这个形式。
            ],
            # messages=[
            #     {"role": "system", "content": system_prompt},
            #     {"role": "user", "content": st.session_state.messages[0]['content']},
            # ], ### original code here.
            stream=False,
        )
        
            llm_response = chatgpt_response.choices[0].message.content
            message_placeholder.markdown('结果如下：')
            
        st.session_state.messages.append(
        {"role": "assistant", "content": llm_response})
    
    # llm_response = chatgpt.chatgpt(user_prompt=final_prompt) ### original code here.
    
    ###NOTE: 需要去除mermaid关键字，需要去除反引号。
    try:
        ### 删除所有的非Mermaid格式的字符，反引号前后的文字。
        # pattern = r"mermaid\n(.*?\n?)"
        # match = re.search(pattern, llm_response, re.DOTALL)
        # if match:
        #     md_input = match.group(1)
        # md_input = str(re.split(r'^(.?)\nmermaid\n(.*?)\n*\n*(.*?)$', llm_response, flags=re.MULTILINE|re.DOTALL))
        
        
        ### 尝试各种正则式，为了只保留Mermaid Markdown格式的内容。
        pattern = r'```mermaid([\s\S]*?)```'
        match = re.search(pattern, llm_response)

        if match:
            md_input = match.group(1).strip()
        # print(md_input)
        
        if md_input:
            md_input = md_input.replace('mermaid', '') ##! working！！！这里需要！
            # md_input = md_input.replace('mermaid', ' ') ##! working！！！这里需要！
        else:
            md_input = llm_response.replace('mermaid', '') ##! working！！！这里需要！
        md_input = md_input.replace("```", '') ###! working！注意这里需要去掉的是三个反引号(```)，而不是一个。
    except Exception as e:
        print('Error:', e)
        md_input = llm_response
    
    # mermaid(md_input)
    html_file = mermaid(md_input)
    
    # ##保存HTML的方法
    # html_file = mermaid(md_input)
    # # print('type of html_file:', type(html_file))
    # with open('output.html', 'w', encoding='utf-8') as f: ### 可以保存HTML文件。
    #     f.write(html_file)
    
    # from html2image import Html2Image
    
    # css_settings = '''
    # .center {
    # margin: auto;
    # height: 500px;
    # width: 500px;
    # }
    # body {
    #     background-color: lightgrey;
    #     height: 100%;
    #     display: grid;
    # }
    # '''

    # css_settings = "body {background: grey;}"
    
    
    ###TODO：尝试将HTML文件转成图片，目前各种方法都不成功。
    # hti = Html2Image()
    # img_path = hti.screenshot(html_file='./output.html', save_as='html.png', size=(500, 200))
    # img_path = hti.screenshot(html_file=html_file, save_as='html.png', css_str=css_settings, size=(500, 200))
    # print('img path:', img_path)
    
    # diagram_time = get_current_time()
    # st.success(body=f'程序运行完成！当前时间：{diagram_time}。', icon='💯')
    
    # return None
    return html_file

### 测试prompt
# markdown_chart("画一个流程图，描述用户访问网站的流程。你只需要提供Mermaid Markdown格式的代码，不需要任何额外的说明或者解释")



## 给按键设置的CSS，据说可以把button放在一起更加紧密。参见： https://discuss.streamlit.io/t/st-button-in-one-line/25966/6
st.markdown("""
            <style>
                div[data-testid="column"] {
                    width: fit-content !important;
                    flex: unset;
                }
                div[data-testid="column"] * {
                    width: fit-content !important;
                }
            </style>
            """, unsafe_allow_html=True)

## TODO：看看是否可以保持页面内容，在download button之后。
# user_input = st.chat_input("说点什么吧...")
def main(input):
    # user_input = st.text_input(label='输入您的问题', placeholder='给我一个复杂的业务需求分析流程图。', label_visibility='visible')
    if input:
        # st.session_state.messages.append(user_input)
        # with st.status('检索中...', expanded=False, state='running') as status:
        with st.chat_message("user"):
            st.markdown(input) ## 这里只需要输出用户的prompt，而不是total prompt。
            st.session_state.messages.append({"role": "user", "content": input})
        spinner = st.spinner('处理中...请耐心等待')
        with spinner:
            html_file, mermaid_code = markdown_chart(input=input)
            diagram_time = get_current_time()
            st.success(body=f'程序运行完成！当前时间：{diagram_time}。', icon='💯')

            ## 可以直接下载HTML文件。让可以展示所有的相关图片。
            if html_file:
                col1, col2, col3 = st.columns([1, 1, 8])
                with col1:
                    ste.download_button(
                        label="下载图例",
                        data=html_file,
                        file_name='mydiagram.html',
                        # mime='text/markdown',
                    )
                    # st.download_button(
                    #     label="下载上述图例",
                    #     data=html_file,
                    #     file_name='mydiagram.html',
                    #     # mime='text/markdown',
                    # )
                    
                with col2:
                    ste.download_button(
                        label="下载代码",
                        data=mermaid_code,
                        file_name='mydiagram.txt',
                        # mime='text/markdown',
                    )
                    # st.download_button(
                    #     label="下载上述图例",
                    #     data=html_file,
                    #     file_name='mydiagram.html',
                    #     # mime='text/markdown',
                    # )

                ##NOTE：上面的download button的高度与一般的st.button不同。
                # with col2:
                #     # st.button("重新生成答案")
                #     print('st.session_state.messages now', st.session_state.messages)
                #     st.button("重新生成答案", on_click=regenerate)
                #     if re_btn:
                #         regenerate()

        return html_file


# user_input = st.text_input("说点什么吧...")
if __name__ == '__main__':
    html_file = main(input=user_input)
    
    ##! working. 需要先在sidebar上设置re_btn，然后在这里调用regenerate函数。而不是在button里面直接用on-click来触发函数。

    # if re_btn:
    #     regenerate()

    try:
        if re_btn:
            regenerate()
    except Exception as e:
        print('Error:', e)
        pass