import os
# 環境変数からDIFY_BASE_URLとDIFY_API_KEY_MYWORKFLOWを取得
DIFY_BASE_URL = os.environ.get("DIFY_BASE_URL", "")
DIFY_API_KEY_MYWORKFLOW = os.environ.get("DIFY_API_KEY_MYWORKFLOW", "")
import gradio as gr
import requests
from langchain_community.document_loaders import YoutubeLoader, UnstructuredPDFLoader, WebBaseLoader
from langchain_community.document_loaders import OnlinePDFLoader
import json

def run_workflow(message,language='en'):
    try:
        if not message.get('text',''):
            file = message['files'][0]
            # PDFファイルをロードしてテキストを抽出
            loader = UnstructuredPDFLoader(file)
            data = loader.load()
            raw_text = data[0].page_content
        else:
          text_message = message.get('text','')
          if "youtube.com" in text_message or "youtu.be" in text_message:
              # YouTubeの場合
              loader = YoutubeLoader.from_youtube_url(
                  youtube_url=text_message,
                  add_video_info=True,
                  language=[language],
              )
              docs = loader.load()
              raw_text = str(docs)
          elif text_message.endswith(".pdf"):
              # PDFの場合
              loader = OnlinePDFLoader(text_message)
              docs = loader.load()
              raw_text = docs[0].page_content
          else:
              # それ以外の場合
              loader = WebBaseLoader(text_message)
              docs = loader.load()
              raw_text = docs[0].page_content


        # APIリクエストのための入力データを準備
        inputs = {
            "knowledge": raw_text
        }

        yield raw_text, "loading...", {}

        # APIエンドポイントURL
        url = DIFY_BASE_URL + "/workflows/run"

        # APIリクエストのヘッダー
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {DIFY_API_KEY_MYWORKFLOW}"
        }

        # APIリクエストのデータ
        data = {
            "inputs": inputs,
            "query": "",
            "response_mode": "streaming",
            "user": "abc_123",
        }

        # APIにリクエストを送信
        response = requests.post(url, headers=headers, json=data, stream=True)
        response.raise_for_status()

        assistant_message = ""
        outputs = {}

        # APIレスポンスのチャンク処理
        for chunk in response.iter_lines(delimiter=b"\n\n"):
            if chunk:
                chunk_data = chunk.decode("utf-8").strip()
                if chunk_data.startswith("data:"):
                    json_data = chunk_data[6:]  # "data: "を取り除く
                    if json_data:
                        result = json.loads(json_data)
                        if result.get("event") == "text_chunk":
                            answer = result.get("data", "").get("text", "")
                            assistant_message += str(answer)
                            yield raw_text, assistant_message, result.get("data", "")
                        elif result.get("event") == "workflow_finished":
                            outputs = result.get('data', "")
                            yield raw_text, assistant_message, outputs

    except Exception as e:
        error_message = str(e)
        print(f"Error: {error_message}")
        return "error", error_message, {}

# Gradioインターフェイスの設定
iface = gr.Interface(
    fn=run_workflow,
    inputs=[gr.MultimodalTextbox(label="PDFファイルをアップロード", file_types=[".pdf"], interactive=True),
            gr.Dropdown(label="Language",value="ja",choices=["en","en-US", "ja", "fr","de","it"],allow_custom_value=True)],
    outputs=[
        gr.Textbox(label="生テキスト", show_copy_button=True, max_lines=5),
        gr.Markdown(),
        gr.JSON()
    ],
    title="Documnet to Dify Workflow",
    description="このGradioインターフェイスでは、PDFファイル、YouTube動画、Webページを入力として受け取り、Dify APIワークフローを使用して処理することができます。",
    article="""
    © 2024 @tregu0458. All rights reserved.
# PDF、YouTube、Web URLを入力としたDifyワークフロー
このGradioインターフェイスでは、PDFファイル、YouTube動画、Webページを入力として受け取り、Dify APIワークフローを使用して処理することができます。PDFファイルをアップロードするか、YouTube URLを提供するか、Web URLを入力するだけで、ワークフローがテキストコンテンツを抽出し、Dify APIを使用して処理します。
## 特徴
- PDFファイル、YouTube動画、Webページを入力として対応
- `langchain_community.document_loaders`の`UnstructuredPDFLoader`と`OnlinePDFLoader`を使用してPDFファイルからテキストを抽出
- `langchain_community.document_loaders`の`YoutubeLoader`を使用してYouTube動画からテキストを抽出
- `langchain_community.document_loaders`の`WebBaseLoader`を使用してWebページからテキストを抽出
- 抽出されたテキストをDify APIワークフローで処理
- リアルタイム更新のためのストリーミングレスポンス処理
- 生テキスト、処理されたテキスト、JSON出力の表示
## 使用方法
1. 入力フィールドにPDFファイルをアップロードするか、YouTube URLを提供するか、Web URLを入力します。
2. ドロップダウンから希望の言語を選択します（デフォルトは "ja"）。
3. "Submit"ボタンをクリックして処理を開始します。
4. インターフェイスには、入力から抽出された生テキスト、Dify APIから処理されたテキスト、JSON出力が表示されます。
## 使用コンポーネント
- Dify API
- Gradio
- langchain_community.document_loaders
## ワークフローの仕様
### 入力
- Knowledge
### 出力
- Result
- Raw Content
### LLM
- gemini-1.5-flash-latest
## 注意事項
- 処理は非同期的に実行され、ワークフローの進行に合わせてリアルタイムで結果が表示されます。
- `yield`文を使用して、処理の途中経過を表示しながら、最終的な結果を返します。
- エラー処理は、処理中に発生する可能性のある例外をキャッチして表示するように実装されています。
`run_workflow`関数は、入力メッセージを受け取り、指定された言語でDify APIワークフローを使用して処理を実行します。関数内では、入力タイプ（PDFファイル、YouTube URL、Web URL）に応じて適切なローダーを使用してテキストを抽出し、Dify APIにリクエストを送信して処理を行います。処理の進行状況に応じて、`yield`文を使用して中間結果を表示し、最終的な結果を返します。エラーが発生した場合は、例外をキャッチして適切なエラーメッセージを表示します。
    """
)

if __name__ == "__main__":
    iface.queue().launch()