Update app.py

app.py

@@ -808,8 +808,9 @@ def process_example(topic):
     """Handle example prompts."""
     process_input(topic, [])
 
+# ── 수정된 process_input 전체 ────────────────────────────────────────────
 def process_input(prompt: str, uploaded_files):
-    # 사용자의 메시지 기록
+    # 사용자 메시지 저장
     if not any(m["role"] == "user" and m["content"] == prompt
                for m in st.session_state.messages):
         st.session_state.messages.append({"role": "user", "content": prompt})
@@ -818,9 +819,9 @@ def process_input(prompt: str, uploaded_files):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        placeholder = st.empty()
-        message_placeholder = st.empty()
-        full_response = ""
+        placeholder          = st.empty()
+        message_placeholder  = st.empty()
+        full_response        = ""
 
         use_web_search = st.session_state.web_search_enabled
         has_uploaded   = bool(uploaded_files)
@@ -831,13 +832,13 @@ def process_input(prompt: str, uploaded_files):
 
             client = get_openai_client()
 
-            # ── ① 시스템 프롬프트
+            # ① 시스템 프롬프트
             selected_cat = st.session_state.get("category_focus", "(None)")
             if selected_cat == "(None)":
                 selected_cat = None
             sys_prompt = get_idea_system_prompt(selected_category=selected_cat)
 
-            # ── ② (선택) 웹 검색 & 파일
+            # ② (선택) 웹 검색 & 파일
             search_content = None
             if use_web_search:
                 status.update(label="Searching the web…")
@@ -850,27 +851,38 @@ def process_input(prompt: str, uploaded_files):
                 with st.spinner("Processing files…"):
                     file_content = process_uploaded_files(uploaded_files)
 
-            # ── ③ 메시지 구성
+            # ③ 유저 메시지 구성
             user_content = prompt
             if search_content:
                 user_content += "\n\n" + search_content
             if file_content:
                 user_content += "\n\n" + file_content
 
+            # ---- 카테고리/하위항목 컨텍스트 추가 ---------------------------
+            def category_context(sel):
+                if sel:   # 특정 카테고리 선택 시
+                    return json.dumps(
+                        {sel: physical_transformation_categories[sel]},
+                        ensure_ascii=False)
+                # (None) → 키 목록만 전달
+                return "ALL_CATEGORIES: " + ", ".join(
+                    physical_transformation_categories.keys())
+
             api_messages = [
                 {"role": "system", "content": sys_prompt},
+                {"role": "system", "name": "category_db",
+                 "content": category_context(selected_cat)},
                 {"role": "user",   "content": user_content},
             ]
+            # --------------------------------------------------------------
 
-            # ── ④ OpenAI 호출
+            # ④ OpenAI 스트리밍 호출
             status.update(label="Generating ideas…")
             stream = client.chat.completions.create(
                 model="gpt-4.1-mini",
                 messages=api_messages,
-                temperature=1,
-                max_tokens=MAX_TOKENS,
-                top_p=1,
-                stream=True
+                temperature=1, max_tokens=MAX_TOKENS,
+                top_p=1, stream=True
             )
             for chunk in stream:
                 if chunk.choices and chunk.choices[0].delta.content:
@@ -879,69 +891,57 @@ def process_input(prompt: str, uploaded_files):
             message_placeholder.markdown(full_response)
             status.update(label="Ideas created!", state="complete")
 
-            # ── ⑤ 이미지 생성
-# ── ⑤ 이미지 생성  (기존 블록 전부 교체) ─────────────────────────────
+            # ⑤ 이미지 생성 (CCM 헤딩 or 레거시 패턴 모두 지원)
             if st.session_state.generate_image and full_response:
-
-    # 1️⃣ ❶ CCM 스타일: "### 이미지 프롬프트"  헤딩 아래 줄
+                # CCM 헤딩
                 ccm_match = re.search(
                     r"###\s*이미지\s*프롬프트\s*\n+([^\n]+)",
                     full_response, flags=re.IGNORECASE)
-
-    # 2️⃣ ❷ 이전 스타일: 테이블 or "Image Prompt:"
                 legacy_match = None
                 if not ccm_match:
                     legacy_match = re.search(
                         r"\|\s*(?:\*\*)?Image\s+Prompt(?:\*\*)?\s*\|\s*([^|\n]+)",
-                        full_response, flags=re.IGNORECASE)
-                    if not legacy_match:
-                        legacy_match = re.search(
-                            r"(?i)Image\s+Prompt\s*[:\-]\s*([^\n]+)",
-                            full_response)
-
-    # 3️⃣ 최종 프롬프트 추출
+                        full_response, flags=re.IGNORECASE) \
+                        or re.search(r"(?i)Image\s+Prompt\s*[:\-]\s*([^\n]+)",
+                                     full_response)
                 match = ccm_match or legacy_match
                 if match:
                     raw_prompt = re.sub(r"[\r\n`\"'\\]", " ",
                                         match.group(1)).strip()
-
                     with st.spinner("아이디어 이미지 생성 중…"):
                         img, cap = generate_image(raw_prompt)
-
                     if img:
                         st.image(img, caption=f"아이디어 시각화 – {cap}")
-            # 대화 기록에 이미지도 저장
                         st.session_state.messages.append({
                             "role": "assistant",
                             "content": "",
                             "image": img,
                             "image_caption": f"아이디어 시각화 – {cap}"
                         })
-            
 
+# ── 신규 헬퍼: 결과 다운로드·저장·채팅 기록에 추가 ──────────────────────────
+            def write_output(md_text: str, prompt: str):
+                """최종 마크다운을 한 번만 UI에 연결하고, 파일 다운·JSON 자동저장까지 수행."""
+    # ① 대화 기록에 답변 저장
+                st.session_state.messages.append({"role": "assistant", "content": md_text})
 
+    # ② 다운로드 버튼
+                st.subheader("Download This Output")
+                col_md, col_html = st.columns(2)
+                col_md.download_button(
+                    "Markdown", md_text,
+                    file_name=f"{prompt[:30]}.md", mime="text/markdown")
+                col_html.download_button(
+                    "HTML", md_to_html(md_text, prompt[:30]),
+                    file_name=f"{prompt[:30]}.html", mime="text/html")
+    
+    # ③ JSON 자동 저장
+                if st.session_state.auto_save:
+                    fn = f"chat_history_auto_{datetime.now():%Y%m%d_%H%M%S}.json"
+                    with open(fn, "w", encoding="utf-8") as fp:
+                        json.dump(st.session_state.messages, fp,
+                                  ensure_ascii=False, indent=2)
 
-            # ── ⑥ 결과 저장
-            st.session_state.messages.append(
-                {"role": "assistant", "content": full_response})
-
-            # ── ⑦ 다운로드 버튼
-            st.subheader("Download This Output")
-            c1, c2 = st.columns(2)
-            c1.download_button("Markdown", full_response,
-                               file_name=f"{prompt[:30]}.md",
-                               mime="text/markdown")
-            c2.download_button("HTML",
-                               md_to_html(full_response, prompt[:30]),
-                               file_name=f"{prompt[:30]}.html",
-                               mime="text/html")
-
-            # ── ⑧ 자동 JSON 저장
-            if st.session_state.auto_save:
-                fn = f"chat_history_auto_{datetime.now():%Y%m%d_%H%M%S}.json"
-                with open(fn, "w", encoding="utf-8") as fp:
-                    json.dump(st.session_state.messages, fp,
-                              ensure_ascii=False, indent=2)
 
         except Exception as e:
             err = str(e)