app.py

# -*- coding: utf-8 -*-
# 財政部財政資訊中心 江信宗

import gradio as gr
import openai
from zhconv_rs import zhconv
import time
import re
import os

MODELS = [
    "Meta-Llama-3.1-405B-Instruct",
    "Meta-Llama-3.1-70B-Instruct",
    "Meta-Llama-3.1-8B-Instruct"
]

API_BASE = "https://api.sambanova.ai/v1"

def create_client(api_key=None):
    """Creates an OpenAI client instance."""
    if api_key:
        openai.api_key = api_key
    else:
        openai.api_key = os.getenv("YOUR_API_KEY")
    return openai.OpenAI(api_key=openai.api_key, base_url=API_BASE)

def chat_with_ai(message, chat_history, system_prompt):
    """Formats the chat history for the API call."""
    # 初始化訊息列表，首先新增系統提示詞
    messages = [{"role": "system", "content": system_prompt}]
    # 遍歷聊天歷史，將使用者和AI機器人的對話新增到訊息列表中
    for tup in chat_history:
        # 獲取字典的第一個鍵（通常是使用者的訊息）
        first_key = list(tup.keys())[0]
        # 獲取字典的最後一個鍵（通常是AI機器人的回應）
        last_key = list(tup.keys())[-1]
        # 將使用者的訊息新增到messages列表中
        messages.append({"role": "user", "content": tup[first_key]})
        # 將AI機器人的回應新增到messages列表中
        messages.append({"role": "assistant", "content": tup[last_key]})
    # 新增當前使用者的訊息
    messages.append({"role": "user", "content": message})
    return messages

def respond(message, chat_history, model, system_prompt, thinking_budget, api_key):
    """Sends the message to the API and gets the response."""
    # 建立OpenAI客戶端
    client = create_client(api_key)
    # 格式化聊天歷史
    messages = chat_with_ai(message, chat_history, system_prompt.format(budget=thinking_budget))
    # 記錄開始時間
    start_time = time.time()
    try:
        # 呼叫API獲取回應
        completion = client.chat.completions.create(model=model, messages=messages)
        response = completion.choices[0].message.content
        # 計算思考時間
        thinking_time = time.time() - start_time
        return response, thinking_time
    except Exception as e:
        # 捕獲並返回錯誤資訊
        error_message = f"Error: {str(e)}"
        return error_message, time.time() - start_time

def parse_response(response):
    """Parses the response from the API."""
    # 使用正規表示式提取回答部分
    answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
    # 使用正規表示式提取反思部分
    reflection_match = re.search(r'<reflection>(.*?)</reflection>', response, re.DOTALL)
    # 提取答案和反思內容，如果沒有匹配則設爲空字串
    answer = answer_match.group(1).strip() if answer_match else ""
    reflection = reflection_match.group(1).strip() if reflection_match else ""
    # 提取所有步驟
    steps = re.findall(r'<step>(.*?)</step>', response, re.DOTALL)
    # 如果沒有提取到答案，則返回原始回應
    if answer == "":
        return response, "", ""
    return answer, reflection, steps

def generate(message, history, model, system_prompt, thinking_budget, api_key):
    """Generates the chatbot response."""
    # 獲取AI回應和思考時間
    response, thinking_time = respond(message, history, model, system_prompt, thinking_budget, api_key)
    # 如果回應是錯誤資訊，直接返回
    if response.startswith("Error:"):
        return history + [({"role": "system", "content": response},)], ""
    # 解析AI的回應
    answer, reflection, steps = parse_response(response)
    answer = zhconv(answer, "zh-tw")
    reflection = zhconv(reflection, "zh-tw")
    # 初始化訊息列表
    messages = []
    # 新增使用者的輸入
    messages.append({"role": "user", "content": f'<div style="text-align: left;">{message}</div>'})
    # 格式化AI回應的步驟和反思
    formatted_steps = [f"Step {i}：{zhconv(step, 'zh-tw')}" for i, step in enumerate(steps, 1)]
    all_steps = "<br>".join(formatted_steps) + f"<br><br>Reflection：{reflection}"
    # 新增AI的推理過程和思考時間
    messages.append({
        "role": "assistant",
        "content": f'<div style="text-align: left;">{all_steps}</div>',
        "metadata": {"title": f"推理過程時間： {thinking_time:.2f} 秒"}
    })
    # 新增AI的最終答案
    messages.append({"role": "assistant", "content": f"<b>{answer}</b>"})
    # 返回更新後的歷史記錄和空字串（用於清空輸入框）
    return history + messages, ""

DEFAULT_SYSTEM_PROMPT = """You are a helpful assistant in normal conversation.
When given a problem to solve, you are an expert problem-solving assistant.
Your task is to provide a detailed, step-by-step solution to a given question.

Follow these instructions carefully:
1. What are the key elements of the problem? What information is given, and what is being asked? Read the given question carefully and reset counter between <count> and </count> to {budget}
2. How can we break down the problem into smaller, manageable parts? What relationships exist between these components? Generate a detailed, logical step-by-step solution.
3. What are possible solutions or explanations for this problem? Can we think of at least three different approaches? Enclose each step of your solution within <step> and </step> tags.
4. You are allowed to use at most {budget} steps (starting budget),
   keep track of it by counting down within tags <count> </count>,
   STOP GENERATING MORE STEPS when hitting 0, you don't have to use all of them.
5. How can we test each hypothesis? What evidence supports or contradicts each one? Are there any potential biases in our reasoning? Do a self-reflection when you are unsure about how to proceed,
   based on the self-reflection and reward, decides whether you need to return
   to the previous steps.
6. Based on our evaluation, what is the most likely solution or explanation? After completing the solution steps, reorganize and synthesize the steps
   into the final answer within <answer> and </answer> tags.
7. What have we learned from this reasoning process? Are there any areas where our reasoning could be improved? Provide a critical, honest and subjective self-evaluation of your reasoning
   process within <reflection> and </reflection> tags.
8. Assign a quality score to your solution as a float between 0.0 (lowest
   quality) and 1.0 (highest quality), enclosed in <reward> and </reward> tags.

Example format:
<count> [starting budget] </count>
<step> [Content of step 1] </step>
<count> [remaining budget] </count>
<step> [Content of step 2] </step>
<reflection> [Evaluation of the steps so far] </reflection>
<reward> [Float between 0.0 and 1.0] </reward>
<count> [remaining budget] </count>
<step> [Content of step 3 or Content of some previous step] </step>
<count> [remaining budget] </count>
...
<step>  [Content of final step] </step>
<count> [remaining budget] </count>
<answer> [Final Answer] </answer> (must give final answer in this format)
<reflection> [Evaluation of the solution] </reflection>
<reward> [Float between 0.0 and 1.0] </reward>
"""

custom_css = """
.center-aligned {
    text-align: center !important;
    color: #ff4081;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
    margin-bottom: -5px !important;
}
.gr-input, .gr-box, .gr-dropdown {
    border-radius: 10px !important;
    border: 2px solid #ff4081 !important;
    margin: 0 !important;
}
.gr-input:focus, .gr-box:focus, .gr-dropdown:focus {
    border-color: #f50057 !important;
    box-shadow: 0 0 0 2px rgba(245,0,87,0.2) !important;
}
.input-background {
    background-color: #ffe0b2 !important;
    padding: 15px !important;
    border-radius: 10px !important;
    margin: 0 !important;
}
.input-background textarea {
    font-size: 18px !important;
    background-color: #ffffff;
    border: 1px solid #f0f8ff;
    border-radius: 8px;
}
.api-background {
    background-color: #FFCFB3 !important;
    padding: 10px !important;
    border-radius: 10px !important;
    margin: 0 !important;
}
.custom-button {
    border-radius: 10px !important;
    background-color: #333333 !important;
    color: white !important;
    font-weight: bold !important;
    transition: all 0.3s ease !important;
}
.custom-button:hover {
    background-color: #000000 !important;
    transform: scale(1.05);
}
.pink-bg {
    background-color: #ff4081 !important;
    border-radius: 10px !important;
}
.pink-bg label, .pink-bg .label-wrap {
    color: white !important;
}
.pink-bg textarea {
    color: black !important;
}
.user-message .message.user {
    background-color: #FFF4B5 !important;
    border-radius: 10px !important;
    padding: 10px !important;
    margin: -8px 0px -8px -20px !important;
}
.assistant-message .message.bot {
    background-color: #B7E0FF !important;
    border-radius: 10px !important;
    padding: 10px !important;
    margin: -8px 0px -8px -20px !important;
}
"""

with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css) as demo:
    gr.Markdown("# 🏹 Multi-Chain Reasoning using Llama-3.1-405B-Instruct. Deployed by 江信宗 🏹", elem_classes="center-aligned")
    chatbot = gr.Chatbot(
        label="Chat",
        show_label=False,
        show_share_button=False,
        show_copy_button=True,
        likeable=True,
        layout="panel",
        type="messages",
        height="auto-max",
        elem_classes=["user-message", "assistant-message"],
        container=True
    )
    msg = gr.Textbox(label="請輸入您的問題：", placeholder="輸入完成後直接按 Enter 開始執行......", autofocus=True, max_lines=10, elem_classes="pink-bg")
    with gr.Row():
        model = gr.Dropdown(choices=MODELS, label="選擇模型", value=MODELS[0], elem_classes="input-background")
        thinking_budget = gr.Slider(minimum=1, maximum=100, value=40, step=1, label="思維規劃", info="模型所能進行的最大思考次數", elem_classes="input-background")
        api_key = gr.Textbox(label="API Key", type="password", placeholder="API authentication key for large language models", elem_classes="api-background")
    clear_button = gr.Button("清除聊天記錄", elem_classes="custom-button")
    clear_button.click(
        lambda: ([], "", gr.Info("已成功清除聊天記錄，歡迎繼續提問....")),
        inputs=None,
        outputs=[chatbot, msg]
    )
    system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, visible=False)
    msg.submit(generate, inputs=[msg, chatbot, model, system_prompt, thinking_budget, api_key], outputs=[chatbot, msg])

demo.load(js=None)

if __name__ == "__main__":
    if "SPACE_ID" in os.environ:
        demo.launch()
    else:
        demo.launch(share=True, show_api=False)