Merge pull request #6 from koji/feat_add-flux-schnell

.gitignore

@@ -0,0 +1,3 @@
+.env
+__pycache__
+__pycache__/

app.py

@@ -3,28 +3,24 @@ from cerebras.cloud.sdk import Cerebras
 import openai
 import os
 from dotenv import load_dotenv
+import base64 # 画像デコード用に追加
+from io import BytesIO # 画像ダウンロード用に追加
+from together import Together # Together AI SDKを追加
 
 # --- RECIPE_BASE_PROMPT のインポート ---
-# prompt.py が存在し、RECIPE_BASE_PROMPTが定義されていると仮定
 try:
     from prompt import RECIPE_BASE_PROMPT
 except ImportError:
-    # エラー処理: prompt.pyが見つからないか、変数が定義されていない場合
     st.error("Error: 'prompt.py' not found or 'RECIPE_BASE_PROMPT' is not defined within it.")
-    st.stop() # 致命的なエラーなのでアプリを停止
-    # RECIPE_BASE_PROMPT = "You are a helpful recipe assistant." # フォールバックが必要な場合
-    # print("Warning: 'prompt.py' not found or 'RECIPE_BASE_PROMPT' not defined. Using a default system prompt.")
+    st.stop()
 
 # --- 定数と設定 ---
-
-# モデル定義
 models = {
     "llama3.1-8b": {"name": "Llama3.1-8b", "tokens": 8192, "developer": "Meta"},
     "llama-3.3-70b": {"name": "Llama-3.3-70b", "tokens": 8192, "developer": "Meta"}
 }
-
-# Optillm用ベースURL (必要に応じて変更)
 BASE_URL = "http://localhost:8000/v1"
+IMAGE_MODEL = "black-forest-labs/FLUX.1-schnell-Free" # 使用する画像生成モデル
 
 # --- 環境変数読み込み ---
 load_dotenv()
@@ -34,180 +30,252 @@ st.set_page_config(page_icon="🤖", layout="wide", page_title="Recipe Infograph
 
 # --- ヘルパー関数 ---
 def contains_injection_keywords(text):
-    """Checks for basic prompt injection keywords."""
     keywords = ["ignore previous", "ignore instructions", "disregard", "forget your instructions", "act as", "you must", "system prompt:"]
     lower_text = text.lower()
     return any(keyword in lower_text for keyword in keywords)
 
 def icon(emoji: str):
-    """Shows an emoji as a Notion-style page icon."""
     st.write(
         f'<span style="font-size: 78px; line-height: 1">{emoji}</span>',
         unsafe_allow_html=True,
     )
 
+# --- 画像生成関数 ---
+@st.cache_data(show_spinner="Generating image...") # 結果をキャッシュ & スピナー表示
+def generate_image_from_prompt(_together_client, prompt_text):
+    """Generates an image using Together AI and returns image bytes."""
+    try:
+        response = _together_client.images.generate(
+            prompt=prompt_text,
+            model=IMAGE_MODEL,
+            width=1024,
+            height=768, # モデルに合わせて調整が必要な場合あり
+            steps=4,    # モデルに合わせて調整が必要な場合あり
+            n=1,
+            response_format="b64_json",
+            # stop=[] # stopは通常不要
+        )
+        if response.data and response.data[0].b64_json:
+            b64_data = response.data[0].b64_json
+            image_bytes = base64.b64decode(b64_data)
+            return image_bytes
+        else:
+            st.error("Image generation failed: No image data received.")
+            return None
+    except Exception as e:
+        st.error(f"Image generation error: {e}", icon="🚨")
+        return None
+
 # --- UI 表示 ---
-icon("🧠 x 🧑‍🍳") # アイコンを修正
+icon("🧠 x 🧑‍🍳")
 st.title("Recipe Infographic Prompt Generator")
 st.subheader("Simply enter a dish name or recipe to easily generate image prompts for stunning recipe infographics", divider="orange", anchor=False)
 
 # --- APIキーの処理 ---
+# Cerebras API Key
 api_key_from_env = os.getenv("CEREBRAS_API_KEY")
 show_api_key_input = not bool(api_key_from_env)
-api_key = None
+cerebras_api_key = None
+
+# Together AI API Key Check
+together_api_key = os.getenv("TOGETHER_API_KEY")
 
 # --- サイドバーの設定 ---
 with st.sidebar:
     st.title("Settings")
 
+    # Cerebras Key Input
     if show_api_key_input:
         st.markdown("### :red[Enter your Cerebras API Key below]")
-        api_key_input = st.text_input("Cerebras API Key:", type="password", key="api_key_input_field")
+        api_key_input = st.text_input("Cerebras API Key:", type="password", key="cerebras_api_key_input_field")
         if api_key_input:
-            api_key = api_key_input
+            cerebras_api_key = api_key_input
     else:
-        api_key = api_key_from_env
-        st.success("✓ API Key loaded from environment")
+        cerebras_api_key = api_key_from_env
+        st.success("✓ Cerebras API Key loaded from environment")
 
+    # Together Key Status
+    if not together_api_key:
+         st.warning("TOGETHER_API_KEY environment variable not set. Image generation will not work.", icon="⚠️")
+    else:
+         st.success("✓ Together API Key loaded from environment") # キー自体は表示しない
+
+    # Model selection
     model_option = st.selectbox(
-        "Choose a model:",
+        "Choose a LLM model:", # ラベルを明確化
         options=list(models.keys()),
         format_func=lambda x: models[x]["name"],
         key="model_select"
     )
 
+    # Max tokens slider
     max_tokens_range = models[model_option]["tokens"]
     default_tokens = min(2048, max_tokens_range)
     max_tokens = st.slider(
-        "Max Tokens:",
+        "Max Tokens (LLM):", # ラベルを明確化
         min_value=512,
         max_value=max_tokens_range,
         value=default_tokens,
         step=512,
-        help="Select the maximum number of tokens for the model's response." # helpテキストを修正
+        help="Select the maximum number of tokens for the language model's response."
     )
 
-    use_optillm = st.toggle("Use Optillm", value=False)
+    use_optillm = st.toggle("Use Optillm (for Cerebras)", value=False) # ラベルを明確化
 
 # --- メインアプリケーションロジック ---
 
-# APIキーが最終的に利用可能かチェック (サイドバーの処理後)
-if not api_key:
-    st.markdown("""
-    ## Cerebras API x Streamlit Demo!
-
-    This simple chatbot app demonstrates how to use Cerebras with Streamlit.
+# APIキー(Cerebras)が最終的に利用可能かチェック
+if not cerebras_api_key:
+    # (以前のエラー表示ロジックと同じ)
+    st.markdown("...") # 省略: APIキーがない場合の説明
+    st.stop()
 
-    To get started:
-    """)
-    if show_api_key_input:
-         # サイドバー入力が表示されている場合
-        st.warning("1. :red[Enter your Cerebras API Key in the sidebar.]")
-    else:
-         # 環境変数から読み込むべきだったが、見つからなかった/空だった場合
-        st.error("1. :red[CEREBRAS_API_KEY environment variable not found or empty.] Please set it in your environment (e.g., in a .env file).")
-    st.markdown("2. Configure your settings and start chatting.") # メッセージを少し変更
-    st.stop() # APIキーがない場合はここで停止
-
-# APIキーが利用可能な場合のみクライアントを初期化
+# APIクライアント初期化 (Cerebras & Together)
 try:
+    # Cerebras Client
     if use_optillm:
-        client = openai.OpenAI(
-            base_url=BASE_URL, # Optillmがlocalhostを使用する場合
-            api_key=api_key
-        )
+        llm_client = openai.OpenAI(base_url=BASE_URL, api_key=cerebras_api_key)
     else:
-        # Cerebras SDKがapi_keyだけで初期化可能か確認
-        # SDKのバージョンや使い方によってはendpoint等の追加設定が必要な場合あり
-        client = Cerebras(api_key=api_key)
-    # st.success("API Client Initialized.") # 任意：初期化成功メッセージ
+        llm_client = Cerebras(api_key=cerebras_api_key)
+
+    # Together Client (APIキーがあれば初期化)
+    image_client = None
+    if together_api_key:
+        image_client = Together(api_key=together_api_key) # 明示的にキーを渡すことも可能
+
 except Exception as e:
-    st.error(f"Failed to initialize API client: {str(e)}", icon="🚨")
-    st.stop() # クライアント初期化失敗時も停止
+    st.error(f"Failed to initialize API client(s): {str(e)}", icon="🚨")
+    st.stop()
 
 # --- チャット履歴管理 ---
 if "messages" not in st.session_state:
     st.session_state.messages = []
+if "generated_images" not in st.session_state:
+     st.session_state.generated_images = {} # 画像データをメッセージIDごとに保存 {msg_idx: image_bytes}
 
 if "selected_model" not in st.session_state:
     st.session_state.selected_model = None
 
-# モデルが変更されたら履歴をクリア
+# モデルが変更されたら履歴をクリア (画像履歴もクリアするかは要検討)
 if st.session_state.selected_model != model_option:
     st.session_state.messages = []
+    st.session_state.generated_images = {} # 画像履歴もクリア
     st.session_state.selected_model = model_option
 
-# チャットメッセージを表示
-for message in st.session_state.messages:
-    avatar = '🤖' if message["role"] == "assistant" else '🦔' # アバターを調整 (ユーザーはハリネズミ?)
+# --- チャットメッセージの表示ループ ---
+# このループでは過去のメッセージを表示し、それぞれに画像生成ボタンをつける
+for idx, message in enumerate(st.session_state.messages):
+    avatar = '🤖' if message["role"] == "assistant" else '🦔'
     with st.chat_message(message["role"], avatar=avatar):
         st.markdown(message["content"])
 
-# --- チャット入力と処理 (インデント修正済み) ---
+        # アシスタントのメッセージで、かつ有効な形式の可能性があり、画像クライアントが利用可能な場合
+        if message["role"] == "assistant" and image_client:
+             # 簡単なチェック: 拒否メッセージではないことを確認
+             lower_content = message["content"].lower()
+             is_likely_prompt = "please provide a valid food dish name" not in lower_content
+
+             if is_likely_prompt:
+                 button_key = f"gen_img_{idx}"
+                 if st.button("Generate Image ✨", key=button_key):
+                     # 画像生成関数を呼び出し、結果をセッション状態に保存
+                     image_bytes = generate_image_from_prompt(image_client, message["content"])
+                     if image_bytes:
+                         st.session_state.generated_images[idx] = image_bytes
+                     # ボタンが押されたら再実行されるので、画像表示は下のブロックで行う
+
+                 # 対応する画像データがセッション状態にあれば表示・ダウンロードボタンを表示
+                 if idx in st.session_state.generated_images:
+                     img_bytes = st.session_state.generated_images[idx]
+                     st.image(img_bytes, caption=f"Generated Image for Prompt #{idx+1}")
+                     st.download_button(
+                         label="Download Image 💾",
+                         data=img_bytes,
+                         file_name=f"recipe_infographic_{idx+1}.png",
+                         mime="image/png",
+                         key=f"dl_img_{idx}"
+                     )
+
+# --- チャット入力と新しいメッセージの処理 ---
 if prompt := st.chat_input("Enter food name/food recipe here..."):
-    # ☆★☆ 入力検証 ☆★☆
+    # 入力検証
     if contains_injection_keywords(prompt):
         st.error("Your input seems to contain instructions. Please provide only the dish name or recipe.", icon="🚨")
-    elif len(prompt) > 4000: # 文字数制限は適切に調整してください
+    elif len(prompt) > 4000:
         st.error("Input is too long. Please provide a shorter recipe or dish name.", icon="🚨")
     else:
-        # ↓↓↓ --- 検証をパスした場合の処理 (ここからインデント) --- ↓↓↓
+        # --- 検証をパスした場合の処理 ---
         st.session_state.messages.append({"role": "user", "content": prompt})
 
-        with st.chat_message("user", avatar='🦔'): # ユーザーアバター
+        # ユーザーメッセージを表示
+        with st.chat_message("user", avatar='🦔'):
             st.markdown(prompt)
 
+        # アシスタントの応答を生成・表示
         try:
-            with st.chat_message("assistant", avatar="🤖"): # アシスタントアバター
+            with st.chat_message("assistant", avatar="🤖"):
                 response_placeholder = st.empty()
                 full_response = ""
 
-                # APIに送信するメッセージリストを作成
                 messages_for_api=[
                     {"role": "system", "content": RECIPE_BASE_PROMPT},
-                    {"role": "user", "content": prompt} # 最新のユーザープロンプトのみ
+                    {"role": "user", "content": prompt}
                 ]
-
-                # ストリーミングで応答を取得
                 stream_kwargs = {
-                   "model": model_option,
-                   "messages": messages_for_api,
-                   "max_tokens": max_tokens,
-                   "stream": True,
+                   "model": model_option, "messages": messages_for_api,
+                   "max_tokens": max_tokens, "stream": True,
                 }
-                response_stream = client.chat.completions.create(**stream_kwargs)
+                # LLM Client を使用
+                response_stream = llm_client.chat.completions.create(**stream_kwargs)
 
                 for chunk in response_stream:
                     chunk_content = ""
-                    # API応答の構造に合わせて調整が必要な場合あり
                     if hasattr(chunk, 'choices') and chunk.choices and hasattr(chunk.choices[0], 'delta') and chunk.choices[0].delta and hasattr(chunk.choices[0].delta, 'content'):
                         chunk_content = chunk.choices[0].delta.content or ""
-
                     if chunk_content:
                         full_response += chunk_content
-                        response_placeholder.markdown(full_response + "▌") # カーソル表示
+                        response_placeholder.markdown(full_response + "▌")
 
-                # 最終的な応答を表示（カーソルなし）
+                # 最終応答表示
                 response_placeholder.markdown(full_response)
 
-                # ☆★☆ 出力検証 ☆★☆
+                # --- ここで新しいアシスタントメッセージに対する処理 ---
+                # 応答を履歴に追加 *してから* インデックスを取得
+                st.session_state.messages.append({"role": "assistant", "content": full_response})
+                new_message_idx = len(st.session_state.messages) - 1 # 新しいメッセージのインデックス
+
+                # 出力検証
                 expected_keywords = ["infographic", "step-by-step", "ingredient", "layout", "minimal style"]
                 lower_response = full_response.lower()
-                is_valid_format = any(keyword in lower_response for keyword in expected_keywords)
-                # システムプロンプトで定義した拒否応答の文字列と一致させる
-                is_refusal = "please provide a valid food dish name or recipe for infographic prompt generation" in lower_response
-
-                if not is_valid_format and not is_refusal:
-                    # 期待される形式でもなく、意図した拒否応答でもない場合
-                    st.warning("The generated response might not contain expected keywords or could indicate an issue.", icon="⚠️")
-                elif is_refusal:
-                    # 意図した拒否応答の場合 (infoレベルで表示)
-                     st.info("Input was determined to be invalid or unrelated. Please provide a valid food dish/recipe.") # メッセージを少し調整
-
-                # アシスタントの応答を履歴に追加
-                st.session_state.messages.append(
-                    {"role": "assistant", "content": full_response})
+                is_valid_format_check = any(keyword in lower_response for keyword in expected_keywords)
+                is_refusal_check = "please provide a valid food dish name or recipe for infographic prompt generation" in lower_response
+
+                if not is_valid_format_check and not is_refusal_check:
+                    st.warning("The generated response might not contain expected keywords...", icon="⚠️")
+                elif is_refusal_check:
+                     st.info("Input was determined to be invalid or unrelated...")
+
+                # 画像生成ボタンと表示エリア (新しいメッセージに対して)
+                # 条件: 画像クライアントがあり、拒否応答でない場合
+                if image_client and not is_refusal_check:
+                    button_key = f"gen_img_{new_message_idx}"
+                    if st.button("Generate Image ✨", key=button_key):
+                        image_bytes = generate_image_from_prompt(image_client, full_response)
+                        if image_bytes:
+                            st.session_state.generated_images[new_message_idx] = image_bytes
+                        # 再実行ループで画像表示
+
+                    # 対応する画像データがあれば表示
+                    if new_message_idx in st.session_state.generated_images:
+                         img_bytes = st.session_state.generated_images[new_message_idx]
+                         st.image(img_bytes, caption=f"Generated Image for Prompt #{new_message_idx+1}")
+                         st.download_button(
+                             label="Download Image 💾",
+                             data=img_bytes,
+                             file_name=f"recipe_infographic_{new_message_idx+1}.png",
+                             mime="image/png",
+                             key=f"dl_img_{new_message_idx}"
+                         )
 
         except Exception as e:
             st.error(f"Error generating response: {str(e)}", icon="🚨")
-        # ↑↑↑ --- ここまでが else 節のインデント内 --- ↑↑↑

requirements.txt

@@ -1,3 +1,5 @@
 cerebras_cloud_sdk
 openai
 python-dotenv
+together
+Pillow