Update app.py

app.py

@@ -2,19 +2,27 @@ import os
 import streamlit as st
 import json
 import anthropic
+import requests
+import logging
+from gradio_client import Client
+
+# 로깅 설정
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s')
 
 # API 설정
 api_key = os.environ.get("API_KEY")
 client = anthropic.Anthropic(api_key=api_key)
 
+# 이미지 생성 API URL
+IMAGE_API_URL = "http://211.233.58.201:7896"
+
 # 최대 토큰 수 설정 (Claude-3 Sonnet의 최대 토큰 수)
 MAX_TOKENS = 7999
 
 def get_system_prompt():
-
     return """
-
-
 당신은 전문 블로그 작성 전문가입니다. 모든 블로그 글 작성 요청에 대해 다음의 8단계 프레임워크를 철저히 따르되, 자연스럽고 매력적인 글이 되도록 작성해야 합니다:
 
 독자 연결 단계 1.1. 공감대 형성을 위한 친근한 인사 1.2. 독자의 실제 고민을 반영한 도입 질문 1.3. 주제에 대한 즉각적 관심 유도
@@ -36,10 +44,69 @@ def get_system_prompt():
 작성 시 준수사항 9.1. 글자 수: 1500-2000자 내외 9.2. 문단 길이: 3-4문장 이내 9.3. 시각적 구분: 소제목, 구분선, 번호 목록 활용 9.4. 톤앤매너: 친근하고 전문적인 대화체 9.5. 데이터: 모든 정보의 출처 명시 9.6. 가독성: 명확한 단락 구분과 강조점 사용
 
 이러한 프레임워크를 바탕으로, 요청받은 주제에 대해 체계적이고 매력적인 블로그 포스트를 작성하겠습니다.
-
 """
 
+def test_image_api_connection():
+    """이미지 API 서버 연결 테스트"""
+    try:
+        client = Client(IMAGE_API_URL)
+        return "이미지 API 연결 성공: 정상 작동 중"
+    except Exception as e:
+        logging.error(f"이미지 API 연결 테스트 실패: {e}")
+        return f"이미지 API 연결 실패: {e}"
 
+def generate_image(prompt, width=768, height=768, guidance=3.5, inference_steps=30, seed=3):
+    """이미지 생성 함수"""
+    if not prompt:
+        return None, "오류: 프롬프트를 입력해주세요"
+    
+    try:
+        client = Client(IMAGE_API_URL)
+        result = client.predict(
+            prompt=prompt,
+            width=int(width),
+            height=int(height),
+            guidance=float(guidance),
+            inference_steps=int(inference_steps),
+            seed=int(seed),
+            do_img2img=False,
+            init_image=None,
+            image2image_strength=0.8,
+            resize_img=True,
+            api_name="/generate_image"
+        )
+        logging.info(f"이미지 생성 성공: {result[1]}")
+        return result[0], f"사용된 시드: {result[1]}"
+    except Exception as e:
+        logging.error(f"이미지 생성 실패: {str(e)}")
+        return None, f"오류: {str(e)}"
+
+def extract_image_prompt(blog_content, blog_topic):
+    """블로그 내용에서 이미지 생성을 위한 프롬프트 추출"""
+    image_prompt_system = f"""
+    다음은 '{blog_topic}'에 관한 블로그 글입니다. 이 블로그 글의 내용을 기반으로 적절한 이미지를 생성하기 위한 
+    프롬프트를 작성해주세요. 프롬프트는 영어로 작성하고, 구체적인 시각적 요소를 담아야 합니다.
+    프롬프트만 반환하세요(다른 설명 없이).
+    
+    예시 형식:
+    "A professional photo of [subject], [specific details], [atmosphere], [lighting], [perspective], high quality, detailed"
+    """
+    
+    try:
+        response = client.messages.create(
+            model="claude-3-7-sonnet-20250219",
+            max_tokens=150,
+            system=image_prompt_system,
+            messages=[{"role": "user", "content": blog_content}]
+        )
+        
+        # 응답에서 프롬프트 추출
+        image_prompt = response.content[0].text.strip()
+        logging.info(f"생성된 이미지 프롬프트: {image_prompt}")
+        return image_prompt
+    except Exception as e:
+        logging.error(f"이미지 프롬프트 생성 오류: {e}")
+        return f"A professional photo related to {blog_topic}, detailed, high quality"
 
 def chatbot_interface():
     st.title("Ginigen Blog")
@@ -47,13 +114,22 @@ def chatbot_interface():
     # 모델 고정 설정
     if "ai_model" not in st.session_state:
         st.session_state["ai_model"] = "claude-3-7-sonnet-20250219"
+    
     # 세션 상태 초기화
     if "messages" not in st.session_state:
         st.session_state.messages = []
-        
+    
     # 자동 저장 기능
     if "auto_save" not in st.session_state:
         st.session_state.auto_save = True
+    
+    # 이미지 생성 토글
+    if "generate_image" not in st.session_state:
+        st.session_state.generate_image = False
+    
+    # 이미지 API 상태
+    if "image_api_status" not in st.session_state:
+        st.session_state.image_api_status = test_image_api_connection()
 
     # 대화 기록 관리 (사이드바)
     st.sidebar.title("대화 기록 관리")
@@ -61,6 +137,24 @@ def chatbot_interface():
     # 자동 저장 토글
     st.session_state.auto_save = st.sidebar.toggle("자동 저장", value=st.session_state.auto_save)
     
+    # 이미지 생성 토글
+    st.session_state.generate_image = st.sidebar.toggle("블로그 글 작성 후 이미지 자동 생성", value=st.session_state.generate_image)
+    
+    # 이미지 API 상태 표시
+    st.sidebar.text(st.session_state.image_api_status)
+    
+    # 이미지 생성 설정 (토글이 켜져 있을 때만 표시)
+    if st.session_state.generate_image:
+        st.sidebar.subheader("이미지 생성 설정")
+        width = st.sidebar.slider("너비", 256, 1024, 768, 64)
+        height = st.sidebar.slider("높이", 256, 1024, 768, 64)
+        guidance = st.sidebar.slider("가이던스 스케일", 1.0, 20.0, 3.5, 0.1)
+        inference_steps = st.sidebar.slider("인퍼런스 스텝", 1, 50, 30, 1)
+        seed = st.sidebar.number_input("시드", value=3, min_value=0, step=1)
+    else:
+        # 기본값 설정
+        width, height, guidance, inference_steps, seed = 768, 768, 3.5, 30, 3
+    
     # 대화 기록 불러오기
     uploaded_file = st.sidebar.file_uploader("대화 기록 불러오기", type=['json'])
     if uploaded_file is not None:
@@ -85,6 +179,9 @@ def chatbot_interface():
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
+            # 이미지가 있는 경우 표시
+            if "image" in message:
+                st.image(message["image"], caption=message.get("image_caption", "생성된 이미지"))
 
     # 사용자 입력
     if prompt := st.chat_input("무엇을 도와드릴까요?"):
@@ -110,19 +207,67 @@ def chatbot_interface():
                     
             message_placeholder.markdown(full_response)
             
-        st.session_state.messages.append({"role": "assistant", "content": full_response})
+            # 이미지 생성 옵션이 켜져 있는 경우
+            if st.session_state.generate_image:
+                with st.spinner("블로그에 맞는 이미지 생성 중..."):
+                    # 이미지 프롬프트 생성
+                    image_prompt = extract_image_prompt(full_response, prompt)
+                    
+                    # 이미지 생성
+                    image, image_caption = generate_image(
+                        image_prompt, 
+                        width=width, 
+                        height=height, 
+                        guidance=guidance, 
+                        inference_steps=inference_steps, 
+                        seed=seed
+                    )
+                    
+                    if image:
+                        st.image(image, caption=image_caption)
+                        # 이미지 정보를 응답에 포함
+                        st.session_state.messages.append({
+                            "role": "assistant", 
+                            "content": full_response,
+                            "image": image,
+                            "image_caption": image_caption
+                        })
+                    else:
+                        st.error(f"이미지 생성 실패: {image_caption}")
+                        st.session_state.messages.append({
+                            "role": "assistant", 
+                            "content": full_response
+                        })
+            else:
+                # 이미지 생성 없이 응답만 저장
+                st.session_state.messages.append({
+                    "role": "assistant", 
+                    "content": full_response
+                })
 
         # 자동 저장 기능
         if st.session_state.auto_save:
             try:
+                # 이미지 정보는 저장하지 않음 (JSON에는 바이너리 데이터를 직접 저장할 수 없음)
+                save_messages = []
+                for msg in st.session_state.messages:
+                    save_msg = {"role": msg["role"], "content": msg["content"]}
+                    save_messages.append(save_msg)
+                
                 with open('chat_history_auto_save.json', 'w', encoding='utf-8') as f:
-                    json.dump(st.session_state.messages, f, ensure_ascii=False, indent=4)
+                    json.dump(save_messages, f, ensure_ascii=False, indent=4)
             except Exception as e:
                 st.sidebar.error(f"자동 저장 중 오류 발생: {str(e)}")
 
     # 대화 기록 다운로드
     if st.sidebar.button("대화 기록 다운로드"):
-        json_history = json.dumps(st.session_state.messages, indent=4, ensure_ascii=False)
+        # 이미지 정보는 저장하지 않음
+        save_messages = []
+        for msg in st.session_state.messages:
+            save_msg = {"role": msg["role"], "content": msg["content"]}
+            save_messages.append(save_msg)
+        
+        json_history = json.dumps(save_messages, indent=4, ensure_ascii=False)
         st.sidebar.download_button(
             label="대화 기록 저장하기",
             data=json_history,
@@ -134,4 +279,11 @@ def main():
     chatbot_interface()
 
 if __name__ == "__main__":
+    # requirements.txt 파일 생성
+    with open("requirements.txt", "w") as f:
+        f.write("streamlit>=1.31.0\n")
+        f.write("anthropic>=0.18.1\n")
+        f.write("gradio-client>=1.8.0\n")
+        f.write("requests>=2.32.3\n")
+    
     main()