app.py

import gradio as gr
from huggingface_hub import InferenceClient, HfApi
import os
import requests
from typing import List, Dict, Union
import traceback
from PIL import Image
from io import BytesIO
import asyncio
from gradio_client import Client
import time
import threading
import json

HF_TOKEN = os.getenv("HF_TOKEN")
hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus-08-2024", token=HF_TOKEN)
hf_api = HfApi(token=HF_TOKEN)

def get_headers():
    if not HF_TOKEN:
        raise ValueError("Hugging Face token not found in environment variables")
    return {"Authorization": f"Bearer {HF_TOKEN}"}

def get_space_structure(space_id: str) -> Dict:
    try:
        # space_id에서 owner와 repo_name을 분리
        owner, repo_name = space_id.split('/')
        
        # HfApi를 사용하여 파일 목록을 가져옴
        files = hf_api.list_repo_files(repo_id=space_id, repo_type="space")
        
        # 파일 목록을 트리 구조로 변환
        tree = {"type": "directory", "path": "", "tree": []}
        for file in files:
            path_parts = file.split('/')
            current = tree
            for i, part in enumerate(path_parts):
                if i == len(path_parts) - 1:  # 파일
                    current["tree"].append({"type": "file", "path": part})
                else:  # 디렉토리
                    found = False
                    for item in current["tree"]:
                        if item["type"] == "directory" and item["path"] == part:
                            current = item
                            found = True
                            break
                    if not found:
                        new_dir = {"type": "directory", "path": part, "tree": []}
                        current["tree"].append(new_dir)
                        current = new_dir
        
        return tree
    except Exception as e:
        print(f"Error in get_space_structure: {str(e)}")
        return {"error": f"API request error: {str(e)}"}

def format_tree_structure(tree_data: Dict, indent: str = "") -> str:
    formatted = ""
    for item in tree_data.get("tree", []):
        if item["type"] == "file":
            formatted += f"{indent}├── {item['path']}\n"
        elif item["type"] == "directory":
            formatted += f"{indent}├── {item['path']}/\n"
            formatted += format_tree_structure(item, indent + "│   ")
    return formatted

def format_list_structure(tree_data: Dict) -> List[str]:
    formatted = []
    for item in tree_data.get("tree", []):
        if item["type"] == "file":
            formatted.append(item["path"])
        elif item["type"] == "directory":
            formatted.append(f"{item['path']}/")
            formatted.extend(format_list_structure(item))
    return formatted

def get_app_py_content(space_id: str) -> str:
    app_py_url = f"https://huggingface.co/spaces/{space_id}/raw/main/app.py"
    try:
        response = requests.get(app_py_url, headers=get_headers())
        if response.status_code == 200:
            return response.text
        else:
            return f"app.py file not found or inaccessible for space: {space_id}"
    except requests.RequestException:
        return f"Error fetching app.py content for space: {space_id}"

def summarize_space(space_id: str) -> str:
    system_message = "당신은 Hugging Face Space의 내용을 요약하는 AI 조수입니다. 주어진 정보를 바탕으로 간결하고 명확한 요약을 제공해주세요."
    user_message = f"다음 Hugging Face Space를 요약해주세요: {space_id}"
    
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_message}
    ]
    
    try:
        response = hf_client.chat_completion(messages, max_tokens=400, temperature=0.7)
        return response.choices[0].message.content
    except Exception as e:
        return f"요약 생성 중 오류 발생: {str(e)}"

def take_screenshot(url):
    try:
        print(f"Taking screenshot of URL: {url}")
        client = Client("ginipick/selenium-screenshot-gradio")
        result = client.predict(url=url, api_name="/predict")
        print(f"Screenshot result: {result}")
        if isinstance(result, str) and os.path.exists(result):
            return Image.open(result)
        else:
            print(f"Invalid result from API: {result}")
            return Image.new('RGB', (600, 360), color='lightgray')
    except Exception as e:
        print(f"Screenshot error: {str(e)}")
        return Image.new('RGB', (600, 360), color='lightgray')

def analyze_space(url: str):
    try:
        # URL에서 space_id 추출
        space_id = url.split('spaces/')[-1]
        
        summary = summarize_space(space_id)
        app_content = get_app_py_content(space_id)
        tree_structure = get_space_structure(space_id)
        
        tree_view = format_tree_structure(tree_structure)
        list_view = "\n".join(format_list_structure(tree_structure))
        screenshot = take_screenshot(url)
        
        return summary, app_content, tree_view, list_view, screenshot, url
    except Exception as e:
        print(f"Error in analyze_space: {str(e)}")
        print(traceback.format_exc())
        return f"오류가 발생했습니다: {str(e)}", "", "", "", None, ""

def create_ui():
    try:
        css = """
        footer {visibility: hidden;}
        .minimal-button {min-width: 30px !important; height: 25px !important; line-height: 1 !important; font-size: 12px !important; padding: 2px 5px !important;}
        .space-row {margin-bottom: 5px !important;}
        #refresh-button, #manual-button, #open-space-button {
            width: 100% !important;
            margin-top: 5px !important;
        }
        #info-output, #usage-guide, #tree-view, #list-view {
            height: 400px !important;
            overflow-y: auto !important;
            padding-right: 10px !important;
        }
        #app-py-content {
            height: auto !important;
            max-height: none !important;
            overflow-y: visible !important;
        }
        .output-group {
            border: 1px solid #ddd;
            border-radius: 5px;
            padding: 10px;
            margin-bottom: 20px;
        }
        .scroll-lock {
            overflow: auto !important;
            max-height: 400px !important;
        }
        .full-height {
            height: auto !important;
            max-height: none !important;
        }
        """

        with gr.Blocks(css=css, theme="Nymbo/Nymbo_Theme") as demo:
            gr.Markdown("# HuggingFace Space Analyzer")
            
            with gr.Row():
                url_input = gr.Textbox(label="HuggingFace Space URL")
                analyze_button = gr.Button("분석")
            
            with gr.Row():
                with gr.Column(scale=2):
                    with gr.Tabs():
                        with gr.TabItem("기본 정보"):
                            with gr.Group(elem_classes="output-group scroll-lock"):
                                info_output = gr.Textbox(label="Space 정보 및 요약", elem_id="info-output", lines=20, max_lines=30)
                            
                            screenshot_output = gr.Image(type="pil", label="Live 화면", height=360, width=600)
                            refresh_button = gr.Button("🔄 서비스 화면", elem_id="refresh-button")
                            
                            with gr.Group(elem_classes="output-group full-height"):
                                app_py_content = gr.Code(language="python", label="메인 소스코드", elem_id="app-py-content", lines=None, max_lines=None)
                            
                            open_space_button = gr.Button("선택한 Space 열기", elem_id="open-space-button")
                        
                        with gr.TabItem("코드 구조 분석"):
                            with gr.Group(elem_classes="output-group scroll-lock"):
                                tree_view = gr.Textbox(label="트리 구조", elem_id="tree-view", lines=30, max_lines=50)
                            with gr.Group(elem_classes="output-group scroll-lock"):
                                list_view = gr.Textbox(label="파일 리스트", elem_id="list-view", lines=30, max_lines=50)

            def open_space_in_browser(url):
                if url:
                    import webbrowser
                    webbrowser.open(url)
                    return f"'{url}' 주소가 새 탭에서 열렸습니다."
                return "선택된 Space가 없습니다."

            analyze_button.click(
                analyze_space,
                inputs=[url_input],
                outputs=[info_output, app_py_content, tree_view, list_view, screenshot_output, url_input]
            )

            open_space_button.click(
                open_space_in_browser,
                inputs=[url_input],
                outputs=[gr.Textbox(label="상태 메시지")]
            )

            refresh_button.click(
                take_screenshot,
                inputs=[url_input],
                outputs=[screenshot_output]
            )

        return demo

    except Exception as e:
        print(f"Error in create_ui: {str(e)}")
        print(traceback.format_exc())
        raise

if __name__ == "__main__":
    try:
        demo = create_ui()
        demo.launch()
    except Exception as e:
        print(f"Error in main: {str(e)}")
        print(traceback.format_exc())