app.py

import os
import re
from http import HTTPStatus
from typing import Dict, List, Optional, Tuple
import base64
import anthropic
from functools import partial

import gradio as gr
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

# 파일 상단의 import 문 아래에 추가
def get_image_base64(image_path):
    with open(image_path, "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode()
    return encoded_string
    
YOUR_API_TOKEN = os.getenv('ANTHROPIC_API_KEY')
client = anthropic.Anthropic(api_key=YOUR_API_TOKEN)

class Role:
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"

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

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

def messages_to_history(messages: Messages) -> History:
    assert messages[0]['role'] == 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):
    pattern = r'```html\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>"



DEMO_CACHE = {}

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


with gr.Blocks(css_paths="app.css") as demo:
    history = gr.State([])
    setting = gr.State({
        "system": SystemPrompt,
    })

    def generation_code(query: Optional[str], _setting: Dict[str, str], _history: Optional[History]):
        if query is None:
            query = ''
        if _history is None:
            _history = []
        
        messages = history_to_messages(_history, _setting['system'])
        system_message = messages[0]['content']
        
        claude_messages = [
            {"role": msg["role"] if msg["role"] != "system" else "user", "content": msg["content"]} 
            for msg in messages[1:] + [{'role': Role.USER, 'content': query}]
        ]
        
        try:
            # 중간 상태 표시
            yield [
                "Generating code...",  # code_output
                _history,              # history
                None,                  # sandbox
                gr.update(active_key="loading"),  # state_tab
                gr.update(open=True)   # code_drawer
            ]

            # 스트리밍 응답 사용
            with client.messages.stream(
                model="claude-3-5-sonnet-20241022",
                max_tokens=7800,
                system=system_message,
                messages=claude_messages
            ) as stream:
                collected_content = ""
                for chunk in stream:
                    if chunk.type == "content_block_delta":
                        delta = chunk.delta.text
                        collected_content += delta
                        
                        yield [
                            collected_content,  # code_output
                            _history,          # history
                            None,              # sandbox (아직 완성되지 않았으므로 None)
                            gr.update(active_key="loading"),  # state_tab
                            gr.update(open=True)  # code_drawer
                        ]
                
                # 최종 결과 반환
                _history = messages_to_history([
                    {'role': Role.SYSTEM, 'content': system_message}
                ] + claude_messages + [{
                    'role': Role.ASSISTANT,
                    'content': collected_content
                }])
                
                yield [
                    collected_content,  # code_output
                    _history,          # history
                    send_to_sandbox(remove_code_block(collected_content)),  # sandbox
                    gr.update(active_key="render"),  # state_tab
                    gr.update(open=True)  # code_drawer
                ]
            
        except Exception as e:
            raise ValueError(f'Error calling Claude API: {str(e)}')


    with ms.Application() as app:
        with antd.ConfigProvider():
            # 메인 컨텐츠 영역
            with antd.Row(gutter=[32, 12], elem_classes="main-content") as layout:  # style 대신 elem_classes 사용
                # 좌측 패널
                with antd.Col(span=24, md=8):
                    with antd.Flex(vertical=True, gap="middle", wrap=True):
                        header = gr.HTML(f"""
                                    <div class="left_header">
                                     <img src="data:image/gif;base64,{get_image_base64('mouse.gif')}" width="360px" />
                                     <h1 style="font-size: 20px;">AI 코딩 코파일럿: MOUSE(WEB)</h2>
                                    </div>
                                     """)
                        input = antd.InputTextarea(
                            size="large", 
                            allow_clear=True, 
                            placeholder="Please enter what kind of application you want"
                        )

                        btn = antd.Button("send", type="primary", size="large")
                        clear_btn = antd.Button("clear history", type="default", size="large")

                # 우측 패널
                with antd.Col(span=24, md=16):
                    with ms.Div(elem_classes="right_panel"):
                        # Setting 버튼들을 우측 상단에 배치
                        with antd.Row(justify="end", elem_classes="setting-buttons-container"):
                            with antd.Space(size="small"):
                                settingPromptBtn = antd.Button(
                                    "⚙️ set system Prompt", type="default")
                                codeBtn = antd.Button("🧑‍💻 view code", type="default")
                                historyBtn = antd.Button("📜 history", type="default")
                        
                        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"):
                                empty = antd.Empty(description="empty input", elem_classes="right_content")
                            with antd.Tabs.Item(key="loading"):
                                loading = antd.Spin(True, tip="coding...", size="large", elem_classes="right_content")
                            with antd.Tabs.Item(key="render"):
                                sandbox = gr.HTML(elem_classes="html_content")

            # Examples 섹션을 별도의 컨테이너로 분리
            with antd.Row(justify="center", elem_classes="examples-container"):
                with antd.Col(span=24):
                    antd.Divider("examples")
                    with antd.Row(gutter=[16, 16], justify="start"):
                        with ms.Each(DEMO_LIST):
                            with antd.Col(xs=24, sm=12, md=8):  # 반응형 그리드
                                with antd.Card(hoverable=True, as_item="card") as demoCard:
                                    antd.CardMeta()
                                demoCard.click(demo_card_click, outputs=[input])

            # Modal과 Drawer 컴포넌트들
            with antd.Modal(open=False, title="set system Prompt", width="800px") as system_prompt_modal:
                systemPromptInput = antd.InputTextarea(
                    SystemPrompt, auto_size=True)

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

            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")


            # 버튼 이벤트 핸들러
            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])

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

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

            btn.click(
                generation_code,
                inputs=[input, setting, history],
                outputs=[code_output, history, sandbox, state_tab, code_drawer]
            )
            
            clear_btn.click(clear_history, inputs=[], outputs=[history])


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