Update app.py

app.py

@@ -1,9 +1,14 @@
+
+
 import os
 import gradio as gr
 import random
 import time
 import logging
+from typing import Iterator
+
 import google.generativeai as genai
+from gradio import ChatMessage  # ChatMessage 구조 사용 (Thinking/Response 구분 가능)
 
 logging.basicConfig(
     level=logging.INFO,
@@ -15,10 +20,16 @@ logging.basicConfig(
 )
 logger = logging.getLogger("idea_generator")
 
+# Gemini API 키 설정
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 genai.configure(api_key=GEMINI_API_KEY)
 
+# 사용할 Gemini 2.0 Flash 모델 (Thinking 기능 포함)
+model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-01-21")
 
+##############################################################################
+# 변환 문자열에서 슬래시("/")로 구분된 두 옵션 중 하나 선택
+##############################################################################
 def choose_alternative(transformation):
     if "/" not in transformation:
         return transformation
@@ -40,6 +51,9 @@ def choose_alternative(transformation):
     else:
         return random.choice([left, right])
 
+##############################################################################
+# 카테고리 사전 (일부만 발췌 가능. 여기서는 예시로 3개만 유지)
+##############################################################################
 
 physical_transformation_categories = {
     "공간 이동": [
@@ -181,38 +195,60 @@ physical_transformation_categories = {
         "라이다 센서/감지", "터치 센서/감지", "제스처 센서/감지", "심박 센서/감지", "혈압 센서/감지"
     ]
 }
+##############################################################################
+# 스트리밍용 Gemini API 함수:
+# - 'Thinking' 단계(아이디어 내부 추론)와 최종 'Response' 단계로 구성
+##############################################################################
+def query_gemini_api_stream(prompt: str) -> Iterator[str]:
+    """
+    Gemini 2.0 Flash with 'Thinking' 부분과 'Response' 부분을 
+    분리하여 스트리밍(Chunk)으로 제공한다.
+    """
+    # chat 초기화 (history 없이 단발성 호출)
+    chat = model.start_chat(history=[])
+    response = chat.send_message(prompt, stream=True)
+
+    thought_buffer = ""
+    response_buffer = ""
+    thinking_complete = False
 
-def query_gemini_api(prompt):
-    try:
-        model = genai.GenerativeModel('gemini-2.0-flash-thinking-exp-01-21')
-        response = model.generate_content(prompt)
-
-        try:
-            if hasattr(response, 'text'):
-                return response.text
-            if hasattr(response, 'candidates') and response.candidates:
-                if len(response.candidates) > 0:
-                    candidate = response.candidates[0]
-                    if hasattr(candidate, 'content'):
-                        content = candidate.content
-                        if hasattr(content, 'parts') and content.parts:
-                            if len(content.parts) > 0:
-                                return content.parts[0].text
-            if hasattr(response, 'parts') and response.parts:
-                if len(response.parts) > 0:
-                    return response.parts[0].text
-            return "Unable to generate a response. API response structure is different than expected."
-        except Exception as inner_e:
-            logger.error(f"Error processing response: {inner_e}")
-            return f"An error occurred while processing the response: {str(inner_e)}"
-    except Exception as e:
-        logger.error(f"Error calling Gemini API: {e}")
-        if "API key not valid" in str(e):
-            return "API key is not valid. Please check your GEMINI_API_KEY environment variable."
-        return f"An error occurred while calling the API: {str(e)}"
-
-
-def enhance_with_llm(base_description, obj_name, category):
+    for chunk in response:
+        # 각 chunk에는 candidates[0].content.parts가 들어있다
+        parts = chunk.candidates[0].content.parts
+        
+        # 예시) parts가 2개이면 (0: Thinking, 1: Response 시작)
+        # 그 외에는 1개씩 끊어서 들어올 수 있음
+        if len(parts) == 2 and not thinking_complete:
+            # 아직 Thinking 중인데, 완성된 Thinking + Response 시작이 한 번에 옴
+            thought_buffer += parts[0].text
+            yield f"[Thinking Chunk] {parts[0].text}"
+            
+            response_buffer = parts[1].text
+            yield f"[Response Start] {parts[1].text}"
+            
+            thinking_complete = True
+        elif thinking_complete:
+            # 이미 Thinking은 끝남 → Response를 이어서 스트리밍
+            current_chunk = parts[0].text
+            response_buffer += current_chunk
+            yield current_chunk
+        else:
+            # Thinking 진행 중 (parts가 1개씩 추가됨)
+            current_chunk = parts[0].text
+            thought_buffer += current_chunk
+            yield f"[Thinking Chunk] {current_chunk}"
+
+    # 스트리밍 ���료 후 최종 결과 한번에 제공할 수도 있음
+    yield f"\n[Final Response]\n{response_buffer}"
+
+##############################################################################
+# 카테고리별 간단 설명을 'Thinking' + 'Response'로 확장 (스트리밍)
+##############################################################################
+def enhance_with_llm_stream(base_description, obj_name, category) -> Iterator[str]:
+    """
+    기존 enhance_with_llm를 스트리밍 형태로 바꾼 함수:
+    'Thinking' + 'Response' 단계를 chunk로 순차 전달
+    """
     prompt = f"""
 다음은 '{obj_name}'의 '{category}' 관련 간단한 설명입니다:
 "{base_description}"
@@ -221,18 +257,24 @@ def enhance_with_llm(base_description, obj_name, category):
 2) 혁신 포인트와 기능성 등을 중심으로 
 3~4문장의 아이디어로 확장해 주세요.
 """
-    return query_gemini_api(prompt)
-
+    # query_gemini_api_stream()로부터 chunk를 받아 그대로 yield
+    for chunk in query_gemini_api_stream(prompt):
+        yield chunk
 
+##############################################################################
+# 한 키워드(오브젝트)에 대한 기본 아이디어(카테고리별) 생성
+##############################################################################
 def generate_single_object_transformations(obj):
     results = {}
     for category, transformations in physical_transformation_categories.items():
         transformation = choose_alternative(random.choice(transformations))
         base_description = f"{obj}이(가) {transformation} 현상을 보인다"
-        results[category] = {"base": base_description, "enhanced": None}
+        results[category] = {"base": base_description, "enhanced": ""}
     return results
 
-
+##############################################################################
+# 2개 키워드 상호작용
+##############################################################################
 def generate_two_objects_interaction(obj1, obj2):
     results = {}
     for category, transformations in physical_transformation_categories.items():
@@ -242,10 +284,12 @@ def generate_two_objects_interaction(obj1, obj2):
             "{obj1}과(와) {obj2}이(가) 충돌하면서 {change}가 일어났다"
         ])
         base_description = template.format(obj1=obj1, obj2=obj2, change=transformation)
-        results[category] = {"base": base_description, "enhanced": None}
+        results[category] = {"base": base_description, "enhanced": ""}
     return results
 
-
+##############################################################################
+# 3개 키워드 상호작용
+##############################################################################
 def generate_three_objects_interaction(obj1, obj2, obj3):
     results = {}
     for category, transformations in physical_transformation_categories.items():
@@ -255,17 +299,12 @@ def generate_three_objects_interaction(obj1, obj2, obj3):
             "{obj1}이(가) {obj2}와(과) {obj3} 사이에서 매개체 역할을 하며 {change}를 촉진했다"
         ])
         base_description = template.format(obj1=obj1, obj2=obj2, obj3=obj3, change=transformation)
-        results[category] = {"base": base_description, "enhanced": None}
-    return results
-
-
-def enhance_descriptions(results, objects):
-    obj_name = " 및 ".join([obj for obj in objects if obj])
-    for category, result in results.items():
-        result["enhanced"] = enhance_with_llm(result["base"], obj_name, category)
+        results[category] = {"base": base_description, "enhanced": ""}
     return results
 
-
+##############################################################################
+# 실제 변환 생성 로직
+##############################################################################
 def generate_transformations(text1, text2=None, text3=None):
     if text2 and text3:
         results = generate_three_objects_interaction(text1, text2, text3)
@@ -276,106 +315,137 @@ def generate_transformations(text1, text2=None, text3=None):
     else:
         results = generate_single_object_transformations(text1)
         objects = [text1]
-    return enhance_descriptions(results, objects)
-
-
-def format_results(results):
-    formatted = ""
-    for category, result in results.items():
-        formatted += f"## {category}\n**기본 아이디어**: {result['base']}\n\n**확장된 아이디어**: {result['enhanced']}\n\n---\n\n"
-    return formatted
-
+    return results, objects
 
 ##############################################################################
-# 스트리밍(Streaming) 방식으로 출력하는 함수: yield를 사용
+# 스트리밍: 각 카테고리별로 'Thinking' + 'Response' 부분을 실시간 전달
 ##############################################################################
-def process_inputs_stream(text1, text2, text3):
-    # 1) 첫 메시지
-    yield "입력값 확인 중..."
+def process_inputs_stream(text1, text2, text3) -> Iterator[list]:
+    """
+    Gradio의 Chatbot 형식에 맞춰서,
+    [(role='assistant'|'user', content=...), ...] 형태로 yield한다.
+    생각(Thinking) 단계와 최종 응답을 분리해서 실시간 전송.
+    """
+    messages = []
+
+    # 1) 입력값 확인
+    yield [("assistant", "입력값 확인 중...")]
     time.sleep(0.3)
-    
+
     text1 = text1.strip() if text1 else None
     text2 = text2.strip() if text2 else None
     text3 = text3.strip() if text3 else None
     if not text1:
-        yield "오류: 최소 하나의 키워드를 입력해주세요."
-        return  # 여기서 함수 종료
-    
-    # 2) 다음 메시지
-    yield "창의적인 모델/컨셉/형상 변화 아이디어 생성 중..."
-    time.sleep(0.3)
-    
-    # 3) 실제 아이디어 생성
-    results = generate_transformations(text1, text2, text3)
-    
-    # 4) 중간 단계 출력
-    yield "결과 포맷팅 중..."
+        yield [("assistant", "오류: 최소 하나의 키워드를 입력해주세요.")]
+        return
+
+    # 2) 아이디어 생성
+    yield [("assistant", "창의적인 모델/컨셉/형상 변화 아이디어 생성 중... (카테고리별 분석)")]
     time.sleep(0.3)
-    
-    # 5) 최종 결과 정리
-    formatted = format_results(results)
-    
-    # 6) 결과 출력
-    yield formatted
-    
-    # 7) 완료
-    yield "완료!"
+    results, objects = generate_transformations(text1, text2, text3)
 
-def get_warning_message():
-    if not GEMINI_API_KEY:
-        return "⚠️ 환경 변수 GEMINI_API_KEY가 설정되지 않았습니다. Gemini API 키를 설정하세요."
-    return ""
+    # 카테고리별 스트리밍 처리
+    obj_name = " 및 ".join([obj for obj in objects if obj])
 
+    for i, (category, result_dict) in enumerate(results.items(), start=1):
+        base_desc = result_dict["base"]
 
-with gr.Blocks(title="키워드 기반 창의적 변화 아이디어 생성기",
+        # 카테고리 안내 출력
+        yield [("assistant", f"**[{i}/{len(results)}] 카테고리:** {category}\n\n기본 아이디어: {base_desc}\n\n지금부터 Thinking + Response를 단계적으로 스트리밍합니다...")]
+        time.sleep(0.5)
+
+        # 스트리밍 LLM 호출
+        thinking_text = ""
+        response_text = ""
+        is_thinking_done = False
+
+        # enhance_with_llm_stream 호출
+        for chunk in enhance_with_llm_stream(base_desc, obj_name, category):
+            if chunk.startswith("[Thinking Chunk]"):
+                # 생각 파트
+                thinking_text += chunk.replace("[Thinking Chunk]", "")
+                messages_to_user = f"**[Thinking]**\n{thinking_text}"
+                yield [("assistant", messages_to_user)]
+            elif chunk.startswith("[Response Start]"):
+                # 응답 시작 시점
+                is_thinking_done = True
+                # 남아있는 부분은 response_text로
+                partial = chunk.replace("[Response Start]", "")
+                response_text += partial
+                messages_to_user = f"**[Final Response 시작]**\n{partial}"
+                yield [("assistant", messages_to_user)]
+            elif chunk.startswith("[Final Response]"):
+                # 최종 종료
+                final = chunk.replace("[Final Response]", "")
+                response_text += f"\n{final}"
+                yield [("assistant", f"**[최종 Response]**\n{response_text.strip()}")]
+            else:
+                # 일반 응답 스트리밍
+                if is_thinking_done:
+                    response_text += chunk
+                    yield [("assistant", f"**[응답 진행]**\n{response_text}") ]
+                else:
+                    thinking_text += chunk
+                    yield [("assistant", f"**[Thinking]**\n{thinking_text}")]
+
+        # 한 카테고리 응답 완료
+        result_dict["enhanced"] = response_text
+
+    # 3) 전체 카테고리 완료
+    yield [("assistant", "**모든 카테고리에 대한 스트리밍이 완료되었습니다!**")]
+
+
+##############################################################################
+# Gradio UI
+##############################################################################
+with gr.Blocks(title="스트리밍 예제: Gemini 2.0 Flash Thinking", 
                theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")) as demo:
-    
-    gr.HTML("""
-    <style>
-    body { background: linear-gradient(135deg, #e0eafc, #cfdef3); font-family: 'Arial', sans-serif; }
-    .gradio-container { padding: 20px; }
-    h1, h2 { text-align: center; }
-    h1 { color: #333; }
-    h2 { color: #555; }
-    .output { background-color: #ffffff; padding: 15px; border-radius: 8px; }
-    .gr-button { background-color: #4CAF50; color: white; border: none; border-radius: 4px; padding: 8px 16px; }
-    </style>
-    """)
-    
-    gr.Markdown("# 🚀 키워드 기반 창의적 변화 아이디어 생성기")
-    gr.Markdown("입력한 **키워드**(최대 3개)를 바탕으로, **창의적인 모델/컨셉/형상 변화**에 대한 이해와 **혁신 포인트**, **기능성** 등을 중심으로 확장된 아이디어를 제시합니다.")
-    
-    warning = gr.Markdown(get_warning_message())
-    
+
+    gr.Markdown("# 🚀 키워드 기반 창의적 변화 아이디어 (Gemini 2.0 Flash Thinking, Streaming)")
+    gr.Markdown("키워드 1~3개를 입력하면, **카테고리별로** 'Thinking'과 'Response'가 실시간 스트리밍됩니다.")
+
+    chatbot = gr.Chatbot(
+        label="카테고리별 아이디어(Thinking + Response) 스트리밍",
+        type="tuple",  # (role, content) 쌍의 리스트로 전달
+        render_markdown=True
+    )
+
     with gr.Row():
         with gr.Column(scale=1):
-            text_input1 = gr.Textbox(label="키워드 1 (필수)", placeholder="예: 스마트폰")
-            text_input2 = gr.Textbox(label="키워드 2 (선택)", placeholder="예: 인공지능")
-            text_input3 = gr.Textbox(label="키워드 3 (선택)", placeholder="예: 헬스케어")
+            text_input1 = gr.Textbox(label="키워드 1 (필수)", placeholder="예: 자동차")
+            text_input2 = gr.Textbox(label="키워드 2 (선택)", placeholder="예: 로봇")
+            text_input3 = gr.Textbox(label="키워드 3 (선택)", placeholder="예: 인공지능")
             submit_button = gr.Button("아이디어 생성하기")
+
+            clear_button = gr.Button("대화 지우기")
         
         with gr.Column(scale=2):
-            with gr.TabItem("창의적인 모델/컨셉/형상 변화 아이디어", id="creative_tab"):
-                # Markdown 출력
-                idea_output = gr.Markdown(label="아이디어 결과")
-
-    gr.Examples(
-        examples=[
-            ["자동차", "", ""],
-            ["스마트폰", "인공지능", ""],
-            ["드론", "인공지능", ""],
-            ["운동화", "웨어러블", "건강"],
-        ],
-        inputs=[text_input1, text_input2, text_input3],
-    )
-    
-    # stream=True 옵션을 통해 함수가 yield하는 문자열을 실시간 출력
+            # 이미 chatbot이 자리를 차지하므로 패스
+            pass
+
+    def clear_chat():
+        return []
+
+    examples = [
+        ["자동차", "", ""],
+        ["스마트폰", "인공지능", ""],
+        ["드론", "인공지능", ""],
+        ["운동화", "웨어러블", "건강"],        
+    ]
+    gr.Examples(examples=examples, inputs=[text_input1, text_input2, text_input3])
+
     submit_button.click(
-        fn=process_inputs_stream, 
+        fn=process_inputs_stream,
         inputs=[text_input1, text_input2, text_input3],
-        outputs=idea_output,
-        stream=True
+        outputs=chatbot,
+        stream=True  # 스트리밍 출력
+    )
+
+    clear_button.click(
+        fn=clear_chat,
+        outputs=chatbot
     )
 
 if __name__ == "__main__":
     demo.launch(debug=True)
+