app.py

import os
import re
import base64
import json
from typing import Dict, List, Optional, Tuple

import gradio as gr
from loguru import logger

import sambanova_gradio
from sambanova_gradio import get_fn

import modelscope_studio.components.base as ms
import modelscope_studio.components.legacy as legacy
import modelscope_studio.components.antd as antd
from config import DEMO_LIST, SystemPrompt

YOUR_API_TOKEN = os.getenv('SAMBANOVA_API_KEY')

History = List[Tuple[str, str]]
Messages = List[Dict[str, str]]

from tree_sitter_languages import get_parser
def check_code_error(code, language='python'):
    language_ext = {'go': 'go', 
                    'html': 'html', 
                    'python': 'python', 
                    'cpp': 'cpp', 
                    # 'sql': 'sql', 
                    'php': 'php', 
                    'javascript': 'javascript', 
                    # 'R': 'r', 
                    'typescript': 'typescript', 
                    'ruby': 'ruby', 
                    'bash': 'bash', 
                    'rust': 'rust', 
                    # 'D': 'd', 
                    'scala': 'scala', 
                    'csharp': 'c_sharp', 
                    'java': 'java',
                    'json': 'json',
                    'julia': 'julia',
                    'css': 'css',
                    'jsdoc': 'jsdoc',
                    'swift': 'swift',
                    'cli': 'cli',
                    }    
    if language not in language_ext.keys():
        return {'result': True, 'error': 'no language in tree_sitter_languages'}
    try:
        parser = get_parser(language_ext[language])
        code_bytes = code.encode('utf-8')
        tree = parser.parse(code_bytes)
        root_node = tree.root_node
        
        visit_list = [root_node]
        ans = []
        while visit_list:
            node = visit_list.pop() 
            if node.type == 'ERROR':
                return {'result': True, 'error': True, 'root_node': root_node}
            is_inner = True if node.children else False
            ans.append([node, is_inner])
            for child in node.children:
                visit_list.append(child)

        return {'result': True, 'error': False, 'root_node': root_node}
    except Exception as e:
        return {'result': False, 'error': f"error:{str(e)}"}

def history_to_messages(history: History, system: str) -> Messages:
    messages = [{'role': 'system', 'content': system}]
    for h in history:
        messages.append({'role': 'user', 'content': h[0]})
        messages.append({'role': 'assistant', 'content': h[1]})
    return messages

def messages_to_history(messages: Messages) -> History:
    assert messages[0]['role'] == 'system'
    history = []
    for q, r in zip(messages[1::2], messages[2::2]):
        history.append([q['content'], r['content']])
    return history

def remove_code_block(text, language='html'):
    pattern = r'```' + language.lower() + '\n(.+?)\n```'
    match = re.search(pattern, text, re.DOTALL)
    if match:
        return match.group(1).strip()
    else:
        return text.strip()

def history_render(history: History):
    return gr.update(open=True), history

def clear_history():
    return []

def send_to_sandbox(code):
    encoded_html = base64.b64encode(code.encode('utf-8')).decode('utf-8')
    data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
    return f"<iframe src=\"{data_uri}\" width=\"100%\" height=\"920px\"></iframe>"

def demo_card_click(e: gr.EventData):
    index = e._data['component']['index']
    return DEMO_LIST[index]['description']


def read_jsonl(fn):
    with open(fn, encoding="utf-8") as f:
        return [json.loads(l) for l in f.readlines()]    

def get_item_by_index(item_list, index, val=None):
    if val is not None and val in item_list:
        return val
    if len(item_list) > index:
        return item_list[index]
    return ""

def get_item_list(locale="zh", programming_language="", category="", difficulty="", debugInfo=""):
    samples = read_jsonl(f'./FullStackBench_data/fsb_{locale}_20241204.jsonl')
    item_list_count = {}
    item_samples = []
    for sample in samples:
        labels = sample['labels']
        content = [sample['content']]
        sample_dict = [labels['programming_language'], f"/{labels['category']}/{labels['difficulty']}/", content]
        if (locale != "" and labels['locale'] == locale): 
            if programming_language == "":
                item = labels['programming_language']
                item_list_count[item] = item_list_count.get(item, 0) + 1
                item_samples.append(sample_dict)
            elif category == "" and labels['programming_language'] == programming_language:
                item = labels['category']
                item_list_count[item] = item_list_count.get(item, 0) + 1
                item_samples.append(sample_dict)
            elif difficulty == "" and labels['programming_language'] == programming_language and labels['category'] == category:
                item = labels['difficulty']
                item_list_count[item] = item_list_count.get(item, 0) + 1
                item_samples.append(sample_dict)
            elif labels['programming_language'] == programming_language and labels['category'] == category and labels['difficulty'] == difficulty:
                item = labels['difficulty']
                item_list_count[item] = item_list_count.get(item, 0) + 1
                item_samples.append(sample_dict)
    item_list = list(item_list_count.keys())
    item_list_dropdown = []
    all_count = 0
    for item in item_list:
        all_count += item_list_count[item]
        item_list_dropdown.append((f"{item}[{item_list_count[item]}]",item))
    item_list_dropdown.insert(0,(f"All[{all_count}]", ""))
    return item_list, item_list_dropdown, item_samples

item_list_locale = ["en", "zh"]
item_list_dropdown_locale = [("en[1687]", "en"), ("zh[1687]", "zh")]
item_list_programming_language, item_list_dropdown_programming_language, item_samples = get_item_list(locale=item_list_locale[0], programming_language="", category="", difficulty="")
item_list_category, item_list_dropdown_category, item_samples = get_item_list(locale=item_list_locale[0], programming_language=get_item_by_index(item_list_programming_language, 0, "python"), category="", difficulty="")
item_list_difficulty, item_list_dropdown_difficulty, item_samples = get_item_list(locale=item_list_locale[0], programming_language=get_item_by_index(item_list_programming_language, 0, "python"), category=get_item_by_index(item_list_category, 0, "Basic Programming"), difficulty="")
item_list_select, item_list_select_dropdown, item_samples = get_item_list(locale=item_list_locale[0], programming_language=get_item_by_index(item_list_programming_language, 0, "python"), category=get_item_by_index(item_list_category, 0, "Basic Programming"), difficulty=get_item_by_index(item_list_difficulty, 0))

def get_input_examples(): 
    input_examples = []
    for demo in DEMO_LIST:
        input_examples.append(['html', demo['title'], demo['description']])
    return input_examples

def code_syntax_check_fun(code, language='html'):
    check_result = check_code_error(code, language=language.lower())
    if check_result['result'] == True:
        if check_result['error'] == True:
            return 'No pass'
        elif check_result['error'] == False:
            return 'Pass'
        else:
            return 'No check [no language in check_languages]'
    else:
        return check_result['error']
    
with gr.Blocks(css_paths="app.css", title="Qwen2.5 Coder Artifacts (+SambaNova+FullStackBench)") as demo:
    history_state = gr.State([])
    setting = gr.State({
        "system": SystemPrompt,
    })

    with gr.Row():
        code_language_example = gr.Textbox(show_label=False, lines=1, container=False, visible=False)
        input_text_title= gr.TextArea(show_label=False, container=False, visible=False)
        input_text_example = gr.TextArea(show_label=False, container=False, visible=False)  
        with gr.Tab("Examples_HTML"): 
            input_examples = get_input_examples()
            gr.Examples(
                    examples=input_examples,
                    inputs=[code_language_example, input_text_title, input_text_example],
                    outputs=None,
                    label='',
                    examples_per_page=10,
                    cache_examples=False,
                )       
        with gr.Tab("Examples_from_FullStackBench"):
            with gr.Row():
                with gr.Column():
                    locale = gr.Dropdown(
                        choices=item_list_dropdown_locale,
                        label="locale",
                        value=get_item_by_index(item_list_locale, 0),
                    ) 
                with gr.Column():
                    programming_language = gr.Dropdown(
                        choices=item_list_dropdown_programming_language,
                        label="programming_language",
                        value=get_item_by_index(item_list_programming_language, 0, "python"),
                    )  
                with gr.Column():
                    category = gr.Dropdown(
                        choices=item_list_dropdown_category,
                        label="category",
                        value=get_item_by_index(item_list_category, 0,  "Basic Programming"),
                    )
                with gr.Column():
                    difficulty = gr.Dropdown(
                        choices=item_list_dropdown_difficulty,
                        label="difficulty",
                        value=get_item_by_index(item_list_difficulty, 0),
                    )
            with gr.Row():
                test_examples = gr.Examples(
                        examples=item_samples,
                        inputs=[code_language_example, input_text_title, input_text_example],
                        outputs=None,
                        label='Examples:',
                        examples_per_page=5,
                        cache_examples=False,
                    )
            def locale_select(locale):
                item_list_programming_language, item_list_dropdown_programming_language, item_samples = get_item_list(locale=locale, programming_language="", category="", difficulty="", debugInfo="locale_1_")
                item_list_category, item_list_dropdown_category, item_samples = get_item_list(locale=locale, programming_language=get_item_by_index(item_list_programming_language, 0), category="", difficulty="", debugInfo="locale_2_")
                item_list_difficulty, item_list_dropdown_difficulty, item_samples = get_item_list(locale=locale, programming_language=get_item_by_index(item_list_programming_language, 0), category=get_item_by_index(item_list_category, 0), difficulty="", debugInfo="locale_3_")
                _, _, item_samples = get_item_list(locale=locale, programming_language=get_item_by_index(item_list_programming_language, 0), category=get_item_by_index(item_list_category, 0), difficulty=get_item_by_index(item_list_difficulty, 0), debugInfo="locale_4_")
                return (gr.Dropdown(choices=item_list_dropdown_programming_language, value=item_list_dropdown_programming_language[0][1]),
                        gr.Dropdown(choices=item_list_dropdown_category, value=item_list_dropdown_category[0][1]),
                        gr.Dropdown(choices=item_list_dropdown_difficulty, value=item_list_dropdown_difficulty[0][1]),
                        gr.Dataset(samples=item_samples))
            locale.select(
                fn=locale_select,
                inputs=[locale],
                outputs=[programming_language, category, difficulty, test_examples.dataset]
            )
            def programming_language_select(locale, programming_language):
                item_list_category, item_list_dropdown_category, item_samples = get_item_list(locale=locale, programming_language=programming_language, category="", difficulty="", debugInfo="programming_language_1_")
                item_list_difficulty, item_list_dropdown_difficulty, item_samples = get_item_list(locale=locale, programming_language=programming_language, category=get_item_by_index(item_list_category, 0), difficulty="", debugInfo="programming_language_2_")
                _, _, item_samples = get_item_list(locale=locale, programming_language=programming_language, category=get_item_by_index(item_list_category, 0), difficulty=get_item_by_index(item_list_difficulty, 0), debugInfo="programming_language_3_")
                return (gr.Dropdown(choices=item_list_dropdown_category, value=item_list_dropdown_category[0][1]),
                        gr.Dropdown(choices=item_list_dropdown_difficulty, value=item_list_dropdown_difficulty[0][1]),
                        gr.Dataset(samples=item_samples))
            programming_language.select(
                fn=programming_language_select,
                inputs=[locale, programming_language],
                outputs=[category, difficulty, test_examples.dataset]
            )                
            def category_select(locale, programming_language, category):
                item_list_difficulty, item_list_dropdown_difficulty, item_samples = get_item_list(locale=locale, programming_language=programming_language, category=category, difficulty="", debugInfo="category_1_")
                _, _, item_samples = get_item_list(locale=locale, programming_language=programming_language, category=category, difficulty=get_item_by_index(item_list_difficulty, 0), debugInfo="category_2_")
                return (gr.Dropdown(choices=item_list_dropdown_difficulty, value=item_list_dropdown_difficulty[0][1]),
                        gr.Dataset(samples=item_samples))
            category.select(
                fn=category_select,
                inputs=[locale, programming_language, category],
                outputs=[difficulty, test_examples.dataset]
            )
            def difficulty_select(locale, programming_language, category, difficulty):
                _, _, item_samples = get_item_list(locale=locale, programming_language=programming_language, category=category, difficulty=difficulty, debugInfo="difficulty_1_")
                return gr.Dataset(samples=item_samples)
            difficulty.select(
                fn=difficulty_select,
                inputs=[locale, programming_language, category, difficulty],
                outputs=[test_examples.dataset]
            ) 

    with ms.Application() as app:
        with antd.ConfigProvider():
            with antd.Row(gutter=[32, 12]) as layout:
                with antd.Col(span=24, md=8):
                    with antd.Flex(vertical=True, gap="middle", wrap=True):
                        header = gr.HTML("""
                            <div style="text-align: center;">
                                <!-- Container for GIFs -->
                                <div style="display: flex; justify-content: center; gap: 20px; margin-bottom: 20px;">
                                    <!-- First GIF -->
                                    <img src="https://huggingface.co/spaces/sambanovasystems/SambaNova-Qwen2.5-Coder-Artifacts/resolve/main/logo.png" width="200px" />
                                </div>
                                <!-- Header Text -->
                                <h1>Qwen2.5-Coder-32B with SambaNova Cloud API</h1>
                            </div>
                        """)

                        input = antd.InputTextarea(size="middle", allow_clear=True, placeholder="Please enter what kind of application you want")
                        code_language = gr.Textbox(label="code language", lines=1)
                        btn = antd.Button("send", type="primary", size="large")                        

                        with antd.Flex(gap="small", wrap=True):
                            settingPromptBtn = antd.Button(
                                "⚙️ set system Prompt", type="default")
                            codeBtn = antd.Button("🧑‍💻 view code", type="default")
                            historyBtn = antd.Button("📜 history", type="default")
                        
                        clear_btn = antd.Button("clear history", type="default", size="large")

                    with antd.Modal(open=False, title="set system Prompt", width="800px") as system_prompt_modal:
                        systemPromptInput = antd.InputTextarea(
                            SystemPrompt, auto_size=True)

                    settingPromptBtn.click(lambda: gr.update(
                        open=True), inputs=[], outputs=[system_prompt_modal])
                    system_prompt_modal.ok(lambda input: ({"system": input}, gr.update(
                        open=False)), inputs=[systemPromptInput], outputs=[setting, system_prompt_modal])
                    system_prompt_modal.cancel(lambda: gr.update(
                        open=False), outputs=[system_prompt_modal])

                    with antd.Drawer(open=False, title="code", placement="left", width="750px") as code_drawer:
                        code_output = legacy.Markdown()

                    codeBtn.click(lambda: gr.update(open=True),
                                  inputs=[], outputs=[code_drawer])
                    code_drawer.close(lambda: gr.update(
                        open=False), inputs=[], outputs=[code_drawer])

                    with antd.Drawer(open=False, title="history", placement="left", width="900px") as history_drawer:
                        history_output = legacy.Chatbot(show_label=False, flushing=False, height=960, elem_classes="history_chatbot")

                    historyBtn.click(history_render, inputs=[history_state], outputs=[history_drawer, history_output])
                    history_drawer.close(lambda: gr.update(
                        open=False), inputs=[], outputs=[history_drawer])

                with antd.Col(span=24, md=16):
                    with ms.Div(elem_classes="right_panel"):
                        gr.HTML('<div class="render_header"><span class="header_btn"></span><span class="header_btn"></span><span class="header_btn"></span></div>')
                        with antd.Tabs(active_key="empty", render_tab_bar="() => null") as state_tab:
                            with antd.Tabs.Item(key="empty", label="Render_empty"):
                                empty = antd.Empty(description="empty input", elem_classes="right_content")
                            with antd.Tabs.Item(key="loading", label="Render_loading"):
                                loading = antd.Spin(True, tip="coding...", size="large", elem_classes="right_content")
                            with antd.Tabs.Item(key="render_html", label="Render_HTML") as render_html:
                                sandbox = gr.HTML(elem_classes="html_content")
                            with antd.Tabs.Item(key="render_other", label="Render_Code") as render_other:
                                sandbox_other = legacy.Markdown()
                    with antd.Row():
                        code_syntax_check = gr.Textbox(label="check code syntax: ", lines=1, interactive=False)

        def generation_code(query: Optional[str], code_language, _setting: Dict[str, str], _history: Optional[History]):
            if query is None:
                query = ''
            if _history is None:
                _history = []

            # Prepare the preprocess and postprocess functions
            def preprocess(message, history):
                if code_language.lower() == 'html':
                    messages = [{'role': 'system', 'content': _setting['system']}]
                else:
                    messages = [{'role': 'system', 'content': ""}]
                for user_msg, assistant_msg in history:
                    messages.append({'role': 'user', 'content': user_msg})
                    messages.append({'role': 'assistant', 'content': assistant_msg})
                messages.append({'role': 'user', 'content': message})
                return {'messages': messages}

            def postprocess(response_text):
                return response_text

            # Get the model from sambanova_gradio
            fn = get_fn(
                model_name='Qwen2.5-Coder-32B-Instruct',
                preprocess=preprocess,
                postprocess=postprocess,
                api_key=YOUR_API_TOKEN
            )

            response_text = ''
            assistant_response = ''
            local_history = _history.copy()

            for content in fn(query, local_history):
                response_text = content
                # Update code_output
                yield {
                    code_output: response_text,
                    state_tab: gr.update(active_key="loading"),
                    code_drawer: gr.update(open=True),
                }

            assistant_response = response_text
            local_history.append([query, assistant_response])

            code = remove_code_block(assistant_response, language=code_language)

            if "```html" in assistant_response:
                sandbox_val = send_to_sandbox(code)
                state_tab_val = gr.update(active_key="render_html")
                render_html_val = gr.update(visible=True)
                render_other_val = gr.update(visible=False)
                sandbox_other_val = ""
            else:
                sandbox_val = ""
                state_tab_val = gr.update(active_key="render_other")
                render_html_val = gr.update(visible=False)
                render_other_val = gr.update(visible=True)
                sandbox_other_val = content

            code_syntax_check_val = code_syntax_check_fun(code, language=code_language)
            yield {
                code_output: assistant_response,
                history_state: local_history,
                sandbox: sandbox_val,
                state_tab: state_tab_val,
                code_drawer: gr.update(open=False),
                render_html: render_html_val,
                render_other: render_other_val,
                sandbox_other: sandbox_other_val,
                code_syntax_check: code_syntax_check_val,
            }

        def code_language_example_change(code_language):
            return code_language            
        code_language_example.change(
                fn=code_language_example_change,
                inputs=[code_language_example], 
                outputs=[code_language], 
            )
        def input_text_example_change(input_text_example):
            return input_text_example            
        input_text_example.change(
                fn=input_text_example_change,
                inputs=[input_text_example], 
                outputs=[input], 
            )
        
        btn.click(generation_code,
                  inputs=[input, code_language, setting, history_state],
                  outputs=[code_output, history_state, sandbox, state_tab, code_drawer,
                           render_html, render_other, sandbox_other, code_syntax_check])

        clear_btn.click(clear_history, inputs=[], outputs=[history_state])

    with gr.Row():
        DESCRIPTION = f'### This demo from [SambaNova-Qwen2.5-Coder-Artifacts](https://huggingface.co/spaces/sambanovasystems/SambaNova-Qwen2.5-Coder-Artifacts). <br>'
        DESCRIPTION += f'More examples from [FullStackBench](https://github.com/bytedance/FullStackBench). <br>'
        DESCRIPTION += f'Thanks for their excellent work.'
        gr.Markdown(DESCRIPTION)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=20).launch(ssr_mode=False, show_api=False)