streamlit_app.py

"""
PLaMoチャットアプリ - Streamlit in Snowflake

このアプリは、plamo on snowflake アプリが提供するcore.completion関数を使用して
チャット機能を実現します。

使用方法:
1. plamo on snowflake アプリを起動:
   CALL <APPLICATION_NAME>.app_public.start_app();
   (APPLICATION_NAMEはPLAMO_2_1_PRIME_FOR_SNOWFLAKEなど)
2. このStreamlitアプリをSnowflake上で作成・実行

注意：
- このアプリは、plamo on snowflake アプリが提供するcore.completion関数を使用してチャット機能を実現します。
- snowflakeの同期的なQuery処理を私用しているため、出力をストリーミングすることはできません。
"""

import streamlit as st
from snowflake.snowpark import Session
from snowflake.snowpark.exceptions import SnowparkSQLException
import json

st.set_page_config(page_title="Talk with PLaMo", page_icon="💬", layout="wide")

# セッション管理
if "messages" not in st.session_state:
    st.session_state.messages = []

# URLクエリパラメータの取得（文字列のみ想定）
def get_query_params() -> dict:
    try:
        return dict(st.query_params)
    except Exception:
        try:
            return st.experimental_get_query_params()
        except Exception:
            return {}


def get_query_param(name: str, default: str = "") -> str:
    params = get_query_params()
    value = params.get(name, default)
    if isinstance(value, list):
        return value[0] if value else default
    return value

# 設定
DEFAULT_TIMEOUT_SECONDS = 120  # デフォルトタイムアウト（秒）
TIMEOUT_SECONDS = st.sidebar.number_input(
    "タイムアウト（秒）",
    min_value=10,
    max_value=600,
    value=DEFAULT_TIMEOUT_SECONDS,
    step=10,
    help="LLM応答のタイムアウト時間を設定します",
)

# データベース名の設定（必須）
application_name_default = get_query_param("application_name", "")
APPLICATION_NAME = st.sidebar.text_input(
    "アプリケーション名（必須）",
    value=application_name_default,
    placeholder="PLAMO_2_1_PRIME_FOR_SNOWFLAKEなど",
    help="core.completion関数が存在するアプリケーション名を入力してください。",
)

# LLMパラメータの設定
st.sidebar.markdown("### LLMパラメータ")

MAX_TOKENS = st.sidebar.slider(
    "Max Tokens",
    min_value=1,
    max_value=4096,
    value=4096,
    step=16,
    help="生成する最大トークン数",
)

TEMPERATURE = st.sidebar.slider(
    "Temperature",
    min_value=0.0,
    max_value=2.0,
    value=0.7,
    step=0.1,
    help="生成のランダム性を制御 (0.0でdeterministic。温度を高くするとランダム性が増加します。)",
)


# Snowflakeセッションの取得
@st.cache_resource
def get_snowflake_session():
    try:
        if "snowflake" in st.secrets:
            return Session.builder.configs(st.secrets["snowflake"]).create()
        else:
            return Session.builder.getOrCreate()
    except Exception:
        return Session.builder.getOrCreate()


# アプリケーション名を取得してタイトルを設定
session_for_app_name = get_snowflake_session()

st.title("💬 Talk with PLaMo")
st.markdown("---")


def escape_sql_string(text: str) -> str:
    return text.replace("'", "''")


def check_application_exists(
    session: Session, application_name: str
) -> tuple[bool, str]:
    """
    アプリケーションが存在するかチェック

    Args:
        session: Snowflakeセッション
        application_name: アプリケーション名

    Returns:
        tuple[bool, str]: (存在するかどうか, エラーメッセージ)
    """
    try:
        # SHOW APPLICATIONSでアプリケーション一覧を取得
        result = session.sql("SHOW APPLICATIONS").collect()
        if result and len(result) > 0:
            app_name_upper = application_name.upper()
            for row in result:
                # 行データからアプリケーション名を取得
                # SHOW APPLICATIONSの結果は通常"name"列を持つ
                row_dict = row.asDict() if hasattr(row, "asDict") else {}
                if "name" in row_dict:
                    if row_dict["name"].upper() == app_name_upper:
                        return True, ""
                # フォールバック: 行の文字列表現から検索
                row_str = str(row).upper()
                if app_name_upper in row_str:
                    return True, ""
        return False, "アプリケーションが見つかりませんでした。"
    except Exception as e:
        return False, f"存在確認に失敗: {str(e)}"


def build_prompt_with_history(
    current_prompt: str, messages: list[dict[str, str]]
) -> str:
    if not messages:
        return current_prompt

    history_parts = []
    for msg in messages[:-1]:
        role = msg.get("role", "")
        content = msg.get("content", "")
        if role == "user":
            history_parts.append(f"ユーザー: {content}")
        elif role == "assistant":
            history_parts.append(f"アシスタント: {content}")

    if history_parts:
        history_text = "\n".join(history_parts)
        return f"{history_text}\nユーザー: {current_prompt}\nアシスタント:"
    else:
        return current_prompt


def call_completion_with_timeout(
    prompt: str,
    session: Session,
    application_name: str,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
    max_tokens: int = 512,
    temperature: float = 0.7,
    messages: list[dict[str, str]] | None = None,
) -> tuple[str, bool]:
    if not application_name or not application_name.strip():
        return "エラー: アプリケーション名が設定されていません。", False

    function_name = f"{application_name.strip()}.core.completion"

    if messages:
        full_prompt = build_prompt_with_history(prompt, messages)
    else:
        full_prompt = prompt

    escaped_prompt = escape_sql_string(full_prompt)

    params = {
        "max_tokens": int(max_tokens),
        "temperature": float(temperature),
    }
    params_json = json.dumps(params, ensure_ascii=False)
    escaped_params = params_json.replace("'", "''")

    sql_query = (
        f"SELECT {function_name}('{escaped_prompt}', "
        f"PARSE_JSON('{escaped_params}')) AS response"
    )

    try:
        try:
            timeout_sql = (
                f"ALTER SESSION SET STATEMENT_TIMEOUT_IN_SECONDS = "
                f"{timeout_seconds}"
            )
            session.sql(timeout_sql).collect()
        except Exception:
            pass

        # クエリを実行
        result = session.sql(sql_query).collect()

        if result and len(result) > 0:
            response_text = result[0]["RESPONSE"] or ""
            if response_text == "Error":
                return "エラー: LLMからの応答がエラーでした", False
            return response_text, True
        else:
            return "エラー: 応答が取得できませんでした", False

    except SnowparkSQLException as e:
        error_code = str(e.error_code) if hasattr(e, "error_code") else ""
        error_message = str(e)

        if (
            "timeout" in error_message.lower()
            or "TIMEOUT" in error_code
            or error_code == "600"
            or "600" in str(e)
        ):
            return (
                f"タイムアウト: {timeout_seconds}秒以内に応答が得られませんでした。"
                "時間を延長するか、プロンプトを短くしてください。",
                False,
            )
        elif (
            "unknown user-defined function" in error_message.lower()
            or "does not exist" in error_message.lower()
            or "not found" in error_message.lower()
            or "42601" in error_code
        ):
            # アプリケーションが存在するかチェック
            exists, check_error = check_application_exists(session, application_name)
            if not exists:
                return (
                    f"エラー: アプリケーション名「{application_name}」が間違っています。\n"
                    "正しいアプリケーション名を入力してください。\n"
                    f"（{check_error}）",
                    False,
                )
            else:
                return (
                    f"エラー: core.completion関数が見つかりません。\n"
                    f"アプリケーション「{application_name}」が正しく"
                    "インストール・起動されているか確認してください。\n"
                    f"（エラー詳細: {error_message}）",
                    False,
                )
        else:
            return f"SQLエラー: {error_message}", False

    except Exception as e:
        error_message = str(e)
        if "timeout" in error_message.lower() or "timed out" in error_message.lower():
            return (
                f"タイムアウト: {timeout_seconds}秒以内に応答が得られませんでした。",
                False,
            )
        else:
            return f"エラーが発生しました: {error_message}", False
    finally:
        try:
            session.sql(
                "ALTER SESSION SET STATEMENT_TIMEOUT_IN_SECONDS = DEFAULT"
            ).collect()
        except Exception:
            pass


application_name_configured = APPLICATION_NAME.strip() if APPLICATION_NAME else ""
if not application_name_configured:
    st.warning(
        "⚠️ アプリケーション名が設定されていません。"
        "サイドバーでアプリケーション名 "
        "(PLAMO_2_1_PRIME_FOR_SNOWFLAKEなど)を入力してください。"
    )

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("メッセージを入力してください..."):
    if not application_name_configured:
        error_msg = (
            "アプリケーション名が設定されていません。"
            "サイドバーでアプリケーション名 (PLAMO_2_1_PRIME_FOR_SNOWFLAKEなど)を入力してください。"
        )
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.session_state.messages.append({"role": "assistant", "content": error_msg})
        with st.chat_message("user"):
            st.markdown(prompt)
        with st.chat_message("assistant"):
            st.error(error_msg)
    else:
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            with st.spinner("考え中..."):
                try:
                    session = get_snowflake_session()

                    timeout = int(TIMEOUT_SECONDS)
                    max_tokens_val = int(MAX_TOKENS)
                    temperature_val = float(TEMPERATURE)
                    response_text, success = call_completion_with_timeout(
                        prompt,
                        session,
                        application_name_configured,
                        timeout,
                        max_tokens_val,
                        temperature_val,
                        st.session_state.messages,
                    )

                    if not success:
                        st.error(response_text)
                    else:
                        st.markdown(response_text)

                    st.session_state.messages.append(
                        {"role": "assistant", "content": response_text}
                    )

                except Exception as e:
                    error_msg = f"予期しないエラーが発生しました: {str(e)}"
                    st.error(error_msg)
                    st.session_state.messages.append(
                        {"role": "assistant", "content": error_msg}
                    )

with st.sidebar:
    st.markdown("---")

    if st.button("チャット履歴をクリア", type="secondary"):
        st.session_state.messages = []
        st.rerun()

    st.markdown("---")
    st.markdown("### 使用方法")
    st.markdown(
        """
    1. plamo on snowflake アプリが起動していることを確認
    2. メッセージを入力して送信
    3. PLaMoが応答を生成します
    """
    )

    st.markdown("---")
    st.markdown("### 注意事項")
    st.markdown(
        """
    - アプリが起動していない場合は、以下を実行してください(起動時は10〜15分のウォームアップ時間が必要です):
      ```sql
      CALL <APPLICATION_NAME>.app_public.start_app();
      (APPLICATION_NAMEはPLAMO_2_1_PRIME_FOR_SNOWFLAKEなど)
      ```
    - アプリケーション名は必須です。core.completion関数が存在するアプリケーション名を指定してください
    - タイムアウトが発生した場合は、タイムアウト時間を延長するか、プロンプトを短くしてください
    """
    )