Update app.py

app.py

@@ -1,26 +1,30 @@
-import gradio as gr
-from huggingface_hub import InferenceClient, HfApi
+
 import os
+import gradio as gr
+from gradio import ChatMessage
+from huggingface_hub import HfApi
 import requests
-from typing import List, Dict, Union, Tuple
+import re
 import traceback
-from PIL import Image
-from io import BytesIO
-import asyncio
-from gradio_client import Client
 import time
 import threading
 import json
-import re
+import asyncio
+from typing import List, Dict, Tuple, Iterator
+import google.generativeai as genai
 
 # HuggingFace 관련 API 키 (스페이스 분석 용)
 HF_TOKEN = os.getenv("HF_TOKEN")
 hf_api = HfApi(token=HF_TOKEN)
 
-# Gemini 2.0 Thinking 모델 관련 API 키 및 클라이언트 (LLM 용)
-G_API_KEY = os.getenv("G_API_KEY")
-gemini_client = InferenceClient("gemini-2.0-flash-thinking-exp-01-21", token=G_API_KEY)
+# Gemini 2.0 Flash Thinking 모델 관련 API 키 및 클라이언트 (LLM 용)
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+genai.configure(api_key=GEMINI_API_KEY)
+model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-01-21")
 
+# --------------------------------------------------
+# 파일 및 스페이스 분석 관련 함수들 (기존 코드 유지)
+# --------------------------------------------------
 def get_headers():
     if not HF_TOKEN:
         raise ValueError("Hugging Face token not found in environment variables")
@@ -75,57 +79,6 @@ def format_tree_structure(tree_data: Dict, indent: str = "") -> str:
             formatted += format_tree_structure(child, indent + "  ")
     return formatted
 
-def summarize_code(app_content: str):
-    system_message = "당신은 Python 코드를 분석하고 요약하는 AI 조수입니다. 주어진 코드를 3줄 이내로 간결하게 요약해주세요."
-    user_message = f"다음 Python 코드를 3줄 이내로 요약해주세요:\n\n{app_content}"
-    
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": user_message}
-    ]
-    
-    try:
-        response = gemini_client.chat_completion(messages, max_tokens=200, temperature=0.7)
-        return response.choices[0].message.content
-    except Exception as e:
-        return f"요약 생성 중 오류 발생: {str(e)}"
-
-def analyze_code(app_content: str):
-    system_message = """당신은 Python 코드를 분석하는 AI 조수입니다. 주어진 코드를 분석하여 다음 항목에 대해 설명해주세요:
-    A. 배경 및 필요성
-    B. 기능적 효용성 및 가치
-    C. 특장점
-    D. 적용 대상 및 타겟
-    E. 기대효과
-    기존 및 유사 프로젝트와 비교하여 분석해주세요. Markdown 형식으로 출력하세요."""
-    user_message = f"다음 Python 코드를 분석해주세요:\n\n{app_content}"
-    
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": user_message}
-    ]
-    
-    try:
-        response = gemini_client.chat_completion(messages, max_tokens=1000, temperature=0.7)
-        return response.choices[0].message.content
-    except Exception as e:
-        return f"분석 생성 중 오류 발생: {str(e)}"
-
-def explain_usage(app_content: str):
-    system_message = "당신은 Python 코드를 분석하여 사용법을 설명하는 AI 조수입니다. 주어진 코드를 바탕으로 마치 화면을 보는 것처럼 사용법을 상세히 설명해주세요. Markdown 형식으로 출력하세요."
-    user_message = f"다음 Python 코드의 사용법을 설명해주세요:\n\n{app_content}"
-    
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": user_message}
-    ]
-    
-    try:
-        response = gemini_client.chat_completion(messages, max_tokens=800, temperature=0.7)
-        return response.choices[0].message.content
-    except Exception as e:
-        return f"사용법 설명 생성 중 오류 발생: {str(e)}"
-
 def adjust_lines_for_code(code_content: str, min_lines: int = 10, max_lines: int = 100) -> int:
     """
     코드 내용에 따라 lines 수를 동적으로 조정합니다.
@@ -138,16 +91,13 @@ def adjust_lines_for_code(code_content: str, min_lines: int = 10, max_lines: int
     Returns:
     - int: 설정된 lines 수
     """
-    # 코드의 줄 수 계산
     num_lines = len(code_content.split('\n'))
-    # 줄 수가 min_lines보다 적다면 min_lines 사용, max_lines보다 크면 max_lines 사용
     return min(max(num_lines, min_lines), max_lines)
 
 def analyze_space(url: str, progress=gr.Progress()):
     try:
         space_id = url.split('spaces/')[-1]
         
-        # Space ID 유효성 검사 수정
         if not re.match(r'^[\w.-]+/[\w.-]+$', space_id):
             raise ValueError(f"Invalid Space ID format: {space_id}")
         
@@ -169,7 +119,6 @@ def analyze_space(url: str, progress=gr.Progress()):
         progress(0.9, desc="사용법 설명 생성 중...")
         usage = explain_usage(app_content)
         
-        # 줄 수 계산하여 lines 설정
         app_py_lines = adjust_lines_for_code(app_content)
         
         progress(1.0, desc="완료")
@@ -179,38 +128,178 @@ def analyze_space(url: str, progress=gr.Progress()):
         print(traceback.format_exc())
         return f"오류가 발생했습니다: {str(e)}", "", None, "", "", "", "", 10
 
-def respond(
-    message: str,
-    history: List[Tuple[str, str]],
-    system_message: str = "",
-    max_tokens: int = 1024,
-    temperature: float = 0.7,
-    top_p: float = 0.9,
-):
-    system_prefix = """당신은 허깅페이스에 특화된 AI 코딩 전문가입니다. 사용자의 질문에 친절하고 상세하게 답변해주세요. 
-    Gradio 특성을 정확히 인식하고 Requirements.txt 누락없이 코딩과 오류를 해결해야 합니다.
-    항상 정확하고 유용한 정보를 제공하도록 노력하세요."""
-    
-    messages = [{"role": "system", "content": f"{system_prefix} {system_message}"}]
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-    for message in gemini_client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.get('content', None)
-        if token:
-            response += token.strip("")
-        yield response
+# --------------------------------------------------
+# Gemini 2.0 Flash Thinking 모델 관련 헬퍼 함수들
+# --------------------------------------------------
+def format_chat_history(messages: list) -> list:
+    """
+    Gradio ChatMessage 객체 리스트를 Gemini가 이해할 수 있는 형식으로 변환합니다.
+    """
+    formatted_history = []
+    for message in messages:
+        # thinking 메시지(메타데이터가 있는 메시지)는 건너뜁니다.
+        if not (hasattr(message, "metadata") and message.metadata):
+            formatted_history.append({
+                "role": "user" if message.role == "user" else "assistant",
+                "parts": [message.content or ""]
+            })
+    return formatted_history
+
+def gemini_chat_completion(system_message: str, user_message: str, max_tokens: int = 200, temperature: float = 0.7) -> str:
+    """
+    시스템 및 유저 메시지를 보내고 스트리밍 응답을 누적하여 최종 텍스트를 반환합니다.
+    """
+    # 초기 대화 기록 구성 (단순 문자열 기반)
+    initial_messages = [
+        ChatMessage(role="system", content=system_message),
+        ChatMessage(role="user", content=user_message)
+    ]
+    chat_history = format_chat_history(initial_messages)
+    chat = model.start_chat(history=chat_history)
+    final_response = ""
+    try:
+        for chunk in chat.send_message(user_message, stream=True):
+            parts = chunk.candidates[0].content.parts
+            # 생각(thinking) 부분과 최종 응답 구분이 있다면 최종 응답을 사용
+            if len(parts) == 2:
+                final_response += parts[1].text
+            else:
+                final_response += parts[0].text
+        return final_response.strip()
+    except Exception as e:
+        return f"LLM 호출 중 오류 발생: {str(e)}"
 
+def summarize_code(app_content: str):
+    system_message = "당신은 Python 코드를 분석하고 요약하는 AI 조수입니다. 주어진 코드를 3줄 이내로 간결하게 요약해주세요."
+    user_message = f"다음 Python 코드를 3줄 이내로 요약해주세요:\n\n{app_content}"
+    try:
+        return gemini_chat_completion(system_message, user_message, max_tokens=200, temperature=0.7)
+    except Exception as e:
+        return f"요약 생성 중 오류 발생: {str(e)}"
+
+def analyze_code(app_content: str):
+    system_message = (
+        "당신은 Python 코드를 분석하는 AI 조수입니다. 주어진 코드를 분석하여 다음 항목에 대해 설명해주세요:\n"
+        "A. 배경 및 필요성\n"
+        "B. 기능적 효용성 및 가치\n"
+        "C. 특장점\n"
+        "D. 적용 대상 및 타겟\n"
+        "E. 기대효과\n"
+        "기존 및 유사 프로젝트와 비교하여 분석해주세요. Markdown 형식으로 출력하세요."
+    )
+    user_message = f"다음 Python 코드를 분석해주세요:\n\n{app_content}"
+    try:
+        return gemini_chat_completion(system_message, user_message, max_tokens=1000, temperature=0.7)
+    except Exception as e:
+        return f"분석 생성 중 오류 발생: {str(e)}"
+
+def explain_usage(app_content: str):
+    system_message = "당신은 Python 코드를 분석하여 사용법을 설명하는 AI 조수입니다. 주어진 코드를 바탕으로 마치 화면을 보는 것처럼 사용법을 상세히 설명해주세요. Markdown 형식으로 출력하세요."
+    user_message = f"다음 Python 코드의 ���용법을 설명해주세요:\n\n{app_content}"
+    try:
+        return gemini_chat_completion(system_message, user_message, max_tokens=800, temperature=0.7)
+    except Exception as e:
+        return f"사용법 설명 생성 중 오류 발생: {str(e)}"
+
+def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
+    """
+    Gemini 2.0 Flash Thinking 모델을 사용하여 스트리밍 응답을 제공합니다.
+    이 함수는 Gradio의 ChatMessage 객체를 사용하며, 실시간 'thinking' 및 최종 응답을 보여줍니다.
+    """
+    if not user_message.strip():
+        messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message. Empty input is not allowed."))
+        yield messages
+        return
+
+    try:
+        print(f"\n=== New Request (Text) ===")
+        print(f"User message: {user_message}")
+
+        # 포맷된 대화 기록 생성
+        chat_history = format_chat_history(messages)
+        chat = model.start_chat(history=chat_history)
+        response = chat.send_message(user_message, stream=True)
+
+        thought_buffer = ""
+        response_buffer = ""
+        thinking_complete = False
+
+        # 초기 'thinking' 메시지 추가
+        messages.append(
+            ChatMessage(
+                role="assistant",
+                content="",
+                metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
+            )
+        )
+
+        for chunk in response:
+            parts = chunk.candidates[0].content.parts
+            current_chunk = parts[0].text
+
+            if len(parts) == 2 and not thinking_complete:
+                thought_buffer += current_chunk
+                print(f"\n=== Complete Thought ===\n{thought_buffer}")
+
+                messages[-1] = ChatMessage(
+                    role="assistant",
+                    content=thought_buffer,
+                    metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
+                )
+                yield messages
+
+                response_buffer = parts[1].text
+                print(f"\n=== Starting Response ===\n{response_buffer}")
+
+                messages.append(
+                    ChatMessage(
+                        role="assistant",
+                        content=response_buffer
+                    )
+                )
+                thinking_complete = True
+
+            elif thinking_complete:
+                response_buffer += current_chunk
+                print(f"\n=== Response Chunk ===\n{current_chunk}")
+
+                messages[-1] = ChatMessage(
+                    role="assistant",
+                    content=response_buffer
+                )
+
+            else:
+                thought_buffer += current_chunk
+                print(f"\n=== Thinking Chunk ===\n{current_chunk}")
+
+                messages[-1] = ChatMessage(
+                    role="assistant",
+                    content=thought_buffer,
+                    metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
+                )
+            yield messages
+
+        print(f"\n=== Final Response ===\n{response_buffer}")
+
+    except Exception as e:
+        print(f"\n=== Error ===\n{str(e)}")
+        messages.append(
+            ChatMessage(
+                role="assistant",
+                content=f"I apologize, but I encountered an error: {str(e)}"
+            )
+        )
+        yield messages
+
+def respond(message: str, history: list) -> Iterator[list]:
+    """
+    기존 respond() 함수는 stream_gemini_response()를 호출하여 대화 형식으로 스트리밍 응답을 제공합니다.
+    """
+    return stream_gemini_response(message, history)
+
+# --------------------------------------------------
+# Gradio UI 구성
+# --------------------------------------------------
 def create_ui():
     try:
         css = """
@@ -292,12 +381,12 @@ def create_ui():
         """
 
         with gr.Blocks(theme="default", css=css) as demo:
-            gr.Markdown("# MOUSE: HF Space Deep-Research", elem_classes="header-markdown")
+            gr.Markdown("# MOUSE Space Analysis", elem_classes="header-markdown")
             
             with gr.Tabs() as tabs:
                 with gr.TabItem("분석"):
                     with gr.Row():
-                        with gr.Column(scale=6):  # 왼쪽 영역
+                        with gr.Column(scale=6):
                             url_input = gr.Textbox(label="HuggingFace Space URL", placeholder="예: https://huggingface.co/spaces/username/space_name")
                             analyze_button = gr.Button("분석", variant="primary")
     
@@ -313,7 +402,7 @@ def create_ui():
                             with gr.Group(elem_classes="output-group tree-view-scroll"):
                                 tree_view_output = gr.Textbox(label="파일 구조 (Tree View)", lines=30)
 
-                        with gr.Column(scale=4):  # 오른쪽 영역
+                        with gr.Column(scale=4):
                             with gr.Group(elem_classes="output-group full-height"):
                                 code_tabs = gr.Tabs()
                                 with code_tabs:
@@ -332,7 +421,10 @@ def create_ui():
                                         )
 
                 with gr.TabItem("AI 코딩"):
-                    chatbot = gr.Chatbot(label="대화", elem_classes="output-group full-height")
+                    chatbot = gr.Chatbot(
+                        label="대화",
+                        elem_classes="output-group full-height"
+                    )
 
                     msg = gr.Textbox(label="메시지", placeholder="메시지를 입력하세요...")
                     
@@ -352,14 +444,9 @@ def create_ui():
                     gr.Examples(examples, inputs=msg)
 
                     def respond_wrapper(message, chat_history, max_tokens, temperature, top_p):
-                        bot_message = ""
-                        for response in respond(message, chat_history, max_tokens=max_tokens, temperature=temperature, top_p=top_p):
-                            bot_message = response  # 마지막 응답 저장
-                            yield "", chat_history + [(message, bot_message)]
+                        # stream_gemini_response가 제너레이터를 반환합니다.
+                        return "", stream_gemini_response(message, chat_history)
     
-                        chat_history.append((message, bot_message))
-                        return "", chat_history
-
                     msg.submit(respond_wrapper, [msg, chatbot, max_tokens, temperature, top_p], [msg, chatbot])
                     
                 with gr.TabItem("Recommended Best"):
@@ -383,7 +470,6 @@ def create_ui():
                 outputs=[requirements_content]
             )
 
-            # 동적으로 app.py의 줄 수 조정
             app_py_content.change(lambda lines: gr.update(lines=lines), inputs=[app_py_content_lines], outputs=[app_py_content])
 
         return demo