Update app.py

app.py

@@ -1,329 +1,308 @@
 import os
-import shutil
-from datetime import datetime
+import tempfile
+from google import genai
+from google.genai import types
 import gradio as gr
-from langchain_chroma import Chroma
-from langchain_community.document_loaders import PyPDFLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_google_genai import ChatGoogleGenerativeAI
-from huggingface_hub import snapshot_download
-
-# ==================================================
-# 設定
-# ==================================================
-
-PAPER_DIR = "papers"
-DB_DIR = "/tmp/chroma_db"  # Read-onlyエラー回避のため、一時ディレクトリを指定
-NOTE_FILE = "research_notes.txt"
-
-os.makedirs(PAPER_DIR, exist_ok=True)
-
-
-# ==================================================
-# Gemini / Embedding モデル設定
-# ==================================================
+import fitz  # PyMuPDF
 
+# =====================================================================
+# 1. 初期設定と環境準備
+# =====================================================================
 API_KEY = os.getenv("GEMINI_API_KEY")
+client = genai.Client(api_key=API_KEY)
 
-if API_KEY is None:
-    raise ValueError(
-        "GEMINI_API_KEY が設定されていません。環境変数を確認してください。"
-    )
-
-# 埋め込みモデルの初期化
-embedding_model = HuggingFaceEmbeddings(
-    model="sentence-transformers/all-MiniLM-L6-v2"
-)
-
-# Geminiモデルの初期化 (安定性の高い2.5-flashを採用)
-llm = ChatGoogleGenerativeAI(
-    model="gemini-2.5-flash",
-    temperature=0.3
-)
-
-# ==================================================
-# PDF → DB 登録・再構築
-# ==================================================
-
-def rebuild_database():
-    pdf_files = [
-        os.path.join(PAPER_DIR, f)
-        for f in os.listdir(PAPER_DIR)
-        if f.endswith(".pdf")
-    ]
-
-    if len(pdf_files) == 0:
-        return "PDFがありません"
-
-    all_docs = []
-
-    for pdf in pdf_files:
-        try:
-            loader = PyPDFLoader(pdf)
-            docs = loader.load()
-            all_docs.extend(docs)
-        except Exception as e:
-            print(f"PDF読み込みエラー ({pdf}): {e}")
-
-    if not all_docs:
-        return "読み込み可能なPDFテキストがありませんでした"
-
-    # テキストの最適化分割
-    splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200
-    )
+# ルート直下を指定
+PDF_DIR = "."
 
-    chunks = splitter.split_documents(all_docs)
-
-    # 既存の古いDBディレクトリを安全に削除して再初期化
-    if os.path.exists(DB_DIR):
-        try:
-            shutil.rmtree(DB_DIR)
-        except Exception as e:
-            print(f"既存DBの削除に失敗しました: {e}")
-
-    # ベクトルDBの構築
-    Chroma.from_documents(
-        documents=chunks,
-        embedding=embedding_model,
-        persist_directory=DB_DIR
-    )
-
-    return f"{len(pdf_files)}件のPDFを登録しました"
-
-
-# ==================================================
-# DB取得
-# ==================================================
-
-def get_db():
-    if not os.path.exists(DB_DIR):
-        return None
-
-    return Chroma(
-        persist_directory=DB_DIR,
-        embedding_function=embedding_model
-    )
-
-
-# ==================================================
-# PDF追加
-# ==================================================
-
-def upload_pdfs(files):
-    if files is None:
-        return "PDF未選択"
-
-    count = 0
-    for file in files:
-        dst = os.path.join(
-            PAPER_DIR,
-            os.path.basename(file.name)
-        )
-        try:
-            shutil.copy(file.name, dst)
-            count += 1
-        except Exception as e:
-            print(f"ファイルコピーエラー: {e}")
-
-    # 追加後にデータベースを再構築
-    msg = rebuild_database()
-
-    return f"{count}件アップロード\n{msg}"
-
-
-# ==================================================
-# 手法検索
-# ==================================================
-
-def search_method(question):
-    db = get_db()
-
-    if db is None:
-        return "PDFを登録してください"
-
-    # 関連ドキュメントの検索
-    docs = db.similarity_search(question, k=5)
-    context = "\n\n".join([d.page_content for d in docs])
-
-    prompt = f"""
-あなたは研究指導教員です。
-
-以下の論文内容を参考に、
-手法や専門用語を初心者にも
-わかりやすく説明してください。
-
-【論文情報】
-{context}
-
-【質問】
-{question}
-
-【回答】
-"""
+knowledge_base = ""
 
+def load_pdf_text(file_path):
+    """PDFファイルからテキストを抽出する関数"""
+    text = ""
     try:
-        response = llm.invoke(prompt)
-        return response.content
+        doc = fitz.open(file_path)
+        for page in doc:
+            text += page.get_text()
+        return text
     except Exception as e:
-        return f"Geminiによる生成エラーが発生しました: {e}"
-
-
-# ==================================================
-# 初回データセット取得 & データベース初期化
-# ==================================================
-
-try:
-    snapshot_download(
-        repo_id="aiTOYAMA/papers",
-        repo_type="dataset",
-        local_dir=PAPER_DIR
-    )
-    print("データセットのダウンロードに成功しました")
-except Exception as e:
-    print(f"データセットのダウンロードに失敗しました (手動での登録が必要です): {e}")
-    
-if not os.path.exists(DB_DIR):
-    print(rebuild_database())
-
-
-# ==================================================
-# 実験検索
-# ==================================================
-
-def search_experiment(question):
-    db = get_db()
-
-    if db is None:
-        return "PDFを登録してください"
-
-    docs = db.similarity_search(question, k=5)
-    context = "\n\n".join([d.page_content for d in docs])
-
-    prompt = f"""
-以下の論文内容を参考に、
-実験設定を整理して回答してください。
-
-抽出項目
-
-・データセット
-・学習率
-・Epoch数
-・評価指標
-・比較手法
-・実験結果
-
-【論文情報】
-{context}
-
-【質問】
-{question}
-
-【回答】
-"""
+        print(f"Error loading {file_path}: {e}")
+        return ""
 
+def update_knowledge_base():
+    """フォルダ内の全てのPDFから知識ベースを更新する関数"""
+    global knowledge_base
     try:
-        response = llm.invoke(prompt)
-        return response.content
-    except Exception as e:
-        return f"Geminiによる生成エラーが発生しました: {e}"
-
-
-# ==================================================
-# 保存
-# ==================================================
-
-def save_note(question, answer):
-    if not question or not answer:
-        return "質問または回答が空のため保存できません"
-
-    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-
+        pdf_files = [f for f in os.listdir(PDF_DIR) if f.lower().endswith('.pdf') and os.path.isfile(os.path.join(PDF_DIR, f))]
+    except:
+        pdf_files = []
+        
+    combined_text = ""
+    file_list_str = "【読み込み済みの論文一覧】\n"
+    
+    if not pdf_files:
+        file_list_str += "・なし (論文PDFをアップロードしてください)\n"
+    
+    for pdf in pdf_files:
+        file_list_str += f"・{pdf}\n"
+        combined_text += f"\n--- 論文: {pdf} ---\n"
+        combined_text += load_pdf_text(os.path.join(PDF_DIR, pdf))
+        
+    knowledge_base = combined_text
+    
+    warning = ""
+    if len(knowledge_base) > 200000:
+        warning = "\n⚠️ 注意: 読み込まれたテキスト量が非常に大きいため、回答の生成に時間がかかったり、API制限に達する可能性があります。"
+        
+    status = f"現在、{len(pdf_files)} 個の論文PDFが読み込まれています。{warning}\n\n{file_list_str}"
+    return status
+
+def clear_database():
+    """データベースフォルダ内のPDFファイルをすべて削除して初期化する関数"""
     try:
-        with open(NOTE_FILE, "a", encoding="utf-8") as f:
-            f.write("\n")
-            f.write("=" * 60)
-            f.write("\n")
-            f.write(f"日時: {now}\n\n")
-            f.write(f"質問:\n{question}\n\n")
-            f.write(f"回答:\n{answer}\n")
-            f.write("\n")
-        return "保存完了"
+        pdf_files = [f for f in os.listdir(PDF_DIR) if f.lower().endswith('.pdf') and os.path.isfile(os.path.join(PDF_DIR, f))]
+        for pdf in pdf_files:
+            os.remove(os.path.join(PDF_DIR, pdf))
     except Exception as e:
-        return f"ファイル保存エラー: {e}"
-
-
-# ==================================================
-# UI (Gradio) ※インターフェース構成は完全再現
-# ==================================================
-
-with gr.Blocks() as demo:
-
-    gr.Markdown("# 研究支援RAGアプリ")
-
-    # タブ1: 論文追加
-    with gr.Tab("論文追加"):
-        pdf_input = gr.File(
-            file_count="multiple",
-            file_types=[".pdf"],
-            label="論文PDF"
-        )
-
-        upload_btn = gr.Button("登録")
-
-        upload_status = gr.Textbox(label="状態")
-
-        upload_btn.click(
-            fn=upload_pdfs,
-            inputs=pdf_input,
-            outputs=upload_status
+        print(f"Error deleting {pdf}: {e}")
+    status = update_knowledge_base()
+    return status, None
+
+init_status = update_knowledge_base()
+
+
+# =====================================================================
+# 2. RAGロジック（generate_content & ストリーミング対応）
+# =====================================================================
+def extract_pure_text(content):
+    """Gradio 6の複雑な内部データから純粋なテキストのみを抽���する関数"""
+    if isinstance(content, list):
+        text_list = []
+        for item in content:
+            if isinstance(item, dict) and "text" in item:
+                if isinstance(item["text"], list):
+                    text_list.append(extract_pure_text(item["text"]))
+                else:
+                    text_list.append(str(item["text"]))
+            elif isinstance(item, dict) and "content" in item:
+                text_list.append(extract_pure_text(item["content"]))
+            elif isinstance(item, str):
+                text_list.append(item)
+        return "".join(text_list)
+    elif isinstance(content, dict):
+        if "text" in content:
+            return extract_pure_text(content["text"])
+        elif "content" in content:
+            return extract_pure_text(content["content"])
+    return str(content)
+
+def rag_chat_stream(message, history, mode):
+    global knowledge_base
+    pure_message = extract_pure_text(message)
+    if not pure_message.strip():
+        yield "ユーザーの質問を正しく読み取れませんでした。もう一度入力してください。"
+        return
+    
+    if mode == "terms":
+        mode_instruction = (
+            "あなたは研究をサポートする優秀なAIアシスタントです。提供された論文資料に基づき、"
+            "専門用語の定義や、提案されている「手法・アルゴリズム・理論」について、背景を含めて数式や箇条書きを用いて分かりやすく解説してください。"
         )
-
-    # タブ2: 手法・専門用語検索
-    with gr.Tab("手法・専門用語検索"):
-        method_question = gr.Textbox(label="質問")
-
-        method_answer = gr.Textbox(
-            label="回答",
-            lines=15
+    else:
+        mode_instruction = (
+            "あなたは研究をサポートする優秀なAIアシスタントです。提供された論文資料に基づき、"
+            "「実験設定、データセット、評価指標、結果、比較グラフの内容、考察」に特化して詳しく分析・回答してください。"
         )
 
-        method_btn = gr.Button("検索")
+    system_instruction = f"""
+{mode_instruction}
+ただし、資料に具体的な記述がない内容に関しては、必ず「提供された論文資料に該当する情報が見つからないため、一般的な知識に基づきGeminiが回答します。」と前置きした上で回答してください。
 
-        method_save = gr.Button("保存")
-
-        method_status = gr.Textbox(label="保存状態")
-
-        method_btn.click(
-            fn=search_method,
-            inputs=method_question,
-            outputs=method_answer
-        )
+===== 信頼できる論文資料 =====
+{knowledge_base}
+============================
+"""
 
-        method_save.click(
-            fn=save_note,
-            inputs=[method_question, method_answer],
-            outputs=method_status
+    gemini_contents = []
+    recent_history = history[-12:] if len(history) > 12 else history
+    
+    for msg in recent_history:
+        pure_text = extract_pure_text(msg["content"])
+        if not pure_text.strip() or "ただいま大変混雑しています" in pure_text:
+            continue
+            
+        role = "user" if msg["role"] == "user" else "model"
+        gemini_contents.append(
+            types.Content(role=role, parts=[types.Part.from_text(text=pure_text)])
         )
+        
+    gemini_contents.append(
+        types.Content(role="user", parts=[types.Part.from_text(text=pure_message)])
+    )
 
-    # タブ3: 実験内容検索
-    with gr.Tab("実験内容検索"):
-        exp_question = gr.Textbox(label="質問")
-
-        exp_answer = gr.Textbox(
-            label="回答",
-            lines=15
+    try:
+        response_stream = client.models.generate_content_stream(
+            model="gemini-2.5-flash",
+            contents=gemini_contents,
+            config=types.GenerateContentConfig(system_instruction=system_instruction)
         )
-
-        exp_btn = gr.Button("検索")
-
-        exp_save = gr.Button("保存")  # 整合性を保つため gr.Button に修正
-
-        exp_status = gr.Textbox(label="保存状態")
-
-        exp_btn.click(
-            fn=search_experiment,
-            inputs=exp_question,
-            outputs=exp_answer
+        
+        current_text = ""
+        for chunk in response_stream:
+            if chunk.text:
+                current_text += chunk.text
+                yield current_text
+                
+    except Exception as e:
+        error_msg = str(e)
+        if "503" in error_msg or "UNAVAILABLE" in error_msg:
+            yield (
+                "⚠️【ただいま大変混雑しています】\n\n"
+                "Gemini APIを提供しているGoogleのサーバー側で一時的にアクセスが集中し、パンクしている状態です。\n"
+                "1〜2分ほど時間を置いてから、もう一度送信ボタンを押してみてください。"
+            )
+        else:
+            yield f"エラーが発生しました: {error_msg}"
+
+
+# =====================================================================
+# 3. アクション関数（ファイル追加・履歴ダウンロード用生成）
+# =====================================================================
+def upload_new_papers(file_objects):
+    if file_objects is None:
+        return "ファイルが選択されていません。", None
+    
+    saved_count = 0
+    for file in file_objects:
+        orig_name = os.path.basename(file.name)
+        dest_path = os.path.join(PDF_DIR, orig_name)
+        with open(file.name, "rb") as src, open(dest_path, "wb") as dst:
+            dst.write(src.read())
+        saved_count += 1
+        
+    status = update_knowledge_base()
+    return status, None
+
+def prepare_terms_download(history):
+    if not history: return None # 🌟非表示切り替えを廃止し、値をクリアするだけに
+    return _generate_file(history, "用語・手法")
+
+def prepare_exp_download(history):
+    if not history: return None
+    return _generate_file(history, "実験内容・結果")
+
+def _generate_file(history, mode_label):
+    """共通のテキストファイル生成ロジック"""
+    custom_filename = f"chat_history_{mode_label}.txt"
+    
+    temp_dir = tempfile.gettempdir()
+    file_path = os.path.join(temp_dir, custom_filename)
+    
+    with open(file_path, "w", encoding="utf-8") as f:
+        f.write(f"==================================================\n")
+        f.write(f" 対話履歴ログ: {mode_label}\n")
+        f.write(f"==================================================\n\n")
+        
+        for msg in history:
+            role_label = "【ユーザー】" if msg["role"] == "user" else "【AIアシスタント】"
+            pure_text = extract_pure_text(msg["content"])
+            f.write(f"{role_label}\n{pure_text}\n")
+            f.write("-" * 50 + "\n\n")
+            
+    return file_path
+
+
+# =====================================================================
+# 4. Gradio UIの構築
+# =====================================================================
+with gr.Blocks(title="研究専用RAGアシスタント") as demo:
+    gr.Markdown("# 🎓 研究専用RAGアシスタント")
+    gr.Markdown("論文PDFの内容をベースに、Geminiがあなたの研究・論文読解をサポートします。")
+    
+    with gr.Accordion("📚 論文データベース管理 (PDF追加・初期化)", open=True):
+        with gr.Row():
+            file_input = gr.File(label="論文PDFをアップロード", file_count="multiple", file_types=[".pdf"])
+            with gr.Column():
+                upload_btn = gr.Button("📂 選択した論文を追加", variant="primary")
+                clear_db_btn = gr.Button("🗑️ データベースを全削除して初期化", variant="stop")
+        status_output = gr.Textbox(value=init_status, label="現在のデータベースステータス", interactive=False, lines=6)
+        
+        upload_btn.click(fn=upload_new_papers, inputs=file_input, outputs=[status_output, file_input])
+        clear_db_btn.click(fn=clear_database, inputs=None, outputs=[status_output, file_input])
+
+    gr.Markdown("### 🔍 論文調査")
+    
+    with gr.Tabs():
+        # タブ1: 専門用語・手法用
+        with gr.TabItem("📖 専門用語・手法を調べる"):
+            gr.Markdown("*論文内の見慣れないキーワードや、提案手法のアルゴリズム、理論について質問してください。*")
+            chatbot_terms = gr.Chatbot(label="対話履歴（用語・手法）", elem_id="terms_chat")
+            msg_terms = gr.Textbox(placeholder="例: この論文で使われている○○という手法のメリットは何ですか？", label="質問を入力")
+            
+            with gr.Row():
+                save_terms_btn = gr.Button("💾 対話履歴を保存 (TXT)", variant="primary")
+                clear_terms = gr.Button("🗑️ 対話履歴をクリア")
+            
+            # 🌟修正：最初から visible=True (常に画面に枠を表示)
+            file_download_terms = gr.File(label="ダウンロード（保存ボタンを押すとここにファイルが出ます）", visible=True)
+            
+            def user_terms(user_message, history):
+                return "", history + [{"role": "user", "content": user_message}]
+            
+            def bot_terms(history):
+                user_msg = history[-1]["content"]
+                history_for_api = history[:-1]
+                history.append({"role": "assistant", "content": ""})
+                
+                for current_response in rag_chat_stream(user_msg, history_for_api, mode="terms"):
+                    history[-1]["content"] = current_response
+                    yield history
+
+            msg_terms.submit(user_terms, [msg_terms, chatbot_terms], [msg_terms, chatbot_terms], queue=False).then(
+                bot_terms, chatbot_terms, chatbot_terms
+            )
+            clear_terms.click(lambda: ([], None), None, [chatbot_terms, file_download_terms], queue=False)
+
+            save_terms_btn.click(
+                fn=prepare_terms_download,
+                inputs=chatbot_terms,
+                outputs=file_download_terms
+            )
+
+        # タブ2: 実験内容用
+        with gr.TabItem("📊 実験内容・結果を調べる"):
+            gr.Markdown("*論文内の実験環境、データセット、グラフの解釈、評価結果について質問してください。*")
+            chatbot_exp = gr.Chatbot(label="対話履歴（実験内容）", elem_id="exp_chat")
+            msg_exp = gr.Textbox(placeholder="例: 実験で使用されたデータセットの規模と、Baselineとの精度の差は？", label="質問を入力")
+            
+            with gr.Row():
+                save_exp_btn = gr.Button("💾 対話履歴を保存 (TXT)", variant="primary")
+                clear_exp = gr.Button("🗑️ 対話履歴をクリア")
+                
+            # 🌟修正：最初から visible=True (常に画面に枠を表示)
+            file_download_exp = gr.File(label="ダウンロード（保存ボタンを押すとここにファイルが出ます）", visible=True)
+            
+            def user_exp(user_message, history):
+                return "", history + [{"role": "user", "content": user_message}]
+            
+            def bot_exp(history):
+                user_msg = history[-1]["content"]
+                history_for_api = history[:-1]
+                history.append({"role": "assistant", "content": ""})
+                
+                for current_response in rag_chat_stream(user_msg, history_for_api, mode="experiments"):
+                    history[-1]["content"] = current_response
+                    yield history
+
+            msg_exp.submit(user_exp, [msg_exp, chatbot_exp], [msg_exp, chatbot_exp], queue=False).then(
+                bot_exp, chatbot_exp, chatbot_exp
+            )
+            clear_exp.click(lambda: ([], None), None, [chatbot_exp, file_download_exp], queue=False)
+
+            save_exp_btn.click(
+                fn=prepare_exp_download,
+                inputs=chatbot_exp,
+                outputs=file_download_exp
+            )
+
+if __name__ == "__main__":
+    demo.launch(theme=gr.themes.Soft())