app.py

## 今後の対応予定
## 1. 複数のファイルの読み込み
## 2. OllamaやLlamaCppなどのローカルモデル対応

import os
import gradio as gr
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.utils import embedding_functions
from openai import OpenAI, OpenAIError
from langchain_text_splitters import RecursiveCharacterTextSplitter
import uuid # For generating unique IDs for chunks

openai_api_key = os.environ.get("OPENAI_API_KEY")
if not openai_api_key:
    raise ValueError("OPENAI_API_KEY環境変数が設定されていません。Hugging Face SpacesのSecretsで設定してください。")

clientai = OpenAI(api_key=openai_api_key)
# embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
embedding_model = SentenceTransformer('pkshatech/GLuCoSE-base-ja')

# ChromaDBのカスタム埋め込み関数
class SBERTEmbeddingFunction(embedding_functions.EmbeddingFunction):
    def __init__(self, model):
        self.model = model
    def __call__(self, texts):
        # sentence-transformersモデルはnumpy配列を返すため、tolist()でPythonリストに変換
        return self.model.encode(texts).tolist()

sbert_ef = SBERTEmbeddingFunction(embedding_model)

# --- ChromaDBクライアントとコレクションの初期化 ---
# インメモリモードで動作させ、新しいPDFがアップロードされるたびにコレクションをリセットします。
# グローバル変数としてクライアントを保持
client = chromadb.Client()
collection_name = "pdf_documents_collection"
# コレクションが存在すれば削除し、新しく作成する（新しいPDFごとにリセットするため）
# （Hugging Face SpacesのインメモリDBはセッションごとにリセットされるため、これは初回起動時のみ意味を持つ）
try:
    client.delete_collection(name=collection_name)
except:
    pass # コレクションが存在しない場合はエラーになるので無視
collection = client.get_or_create_collection(name=collection_name, embedding_function=sbert_ef)

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,      # チャンクの最大文字数
    chunk_overlap=150,   # チャンク間のオーバーラップ文字数
    length_function=len, # 文字数で長さを計算
    separators=["\n\n", "\n", " ", ""] # 分割の優先順位
)

# --- ヘルパー関数 ---
def extract_text_from_pdf(pdf_file_path):
    """PDFファイルからテキストを抽出する"""
    try:
        reader = PdfReader(pdf_file_path)
        text = ""
        for page in reader.pages:
            text += page.extract_text() + "\n"
        return text
    except Exception as e:
        return f"PDFの読み込み中にエラーが発生しました: {e}"

def get_gpt_response(query, context):
    """GPTモデルを使用して質問に回答する"""
    messages = [
        {"role": "system", "content": "あなたは提供されたコンテキストに基づいて質問に答える有益なアシスタントです。コンテキストに情報がない場合は、「提供された情報からは回答できません。」と答えてください。コンテキスト内の情報のみを使用し、自身の知識で補完しないでください。"},
        {"role": "user", "content": f"コンテキスト:\n{context}\n\n質問: {query}\n\n回答:"}
    ]
    try:
        response = clientai.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            max_tokens=900, # 回答の最大トークン数
            temperature=0.5, # 創造性の度合い
        )
        return response.choices[0].message.content.strip()
    except OpenAIError as e:
        return f"OpenAI APIエラー: {e}"
    except Exception as e:
        return f"GPTモデルの呼び出し中に予期せぬエラーが発生しました: {e}"

def upload_pdf_and_process(pdf_files): # 引数を複数ファイル対応に変更
    """複数のPDFファイルをアップロードし、テキストを抽出し、ChromaDBに登録する"""
    if not pdf_files: # リストが空かどうかでチェック
        return "PDFファイルがアップロードされていません。", gr.update(interactive=False)

    processed_files_count = 0
    total_chunks_added = 0
    all_status_messages = []

    # コレクションをクリアする処理は削除。既存のコレクションに追加していく。
    # グローバル変数 collection は既に初期化されているものを使用。

    for pdf_file in pdf_files:
        try:
            pdf_path = pdf_file.name
            file_name = os.path.basename(pdf_path)
            all_status_messages.append(f"PDFファイル '{file_name}' を処理中...")

            # 1. PDFからテキストを抽出
            raw_text = extract_text_from_pdf(pdf_path)
            if "エラー" in raw_text:
                all_status_messages.append(raw_text)
                continue # 次のファイルへ

            # 2. テキストをチャンクに分割
            chunks = text_splitter.split_text(raw_text)
            if not chunks:
                all_status_messages.append(f"'{file_name}' から有効なテキストチャンクを抽出できませんでした。")
                continue # 次のファイルへ

            # 3. チャンクをChromaDBに登録
            documents = chunks
            # メタデータにファイル名を追加
            metadatas = [{"source": file_name, "chunk_index": i} for i in range(len(chunks))]
            ids = [str(uuid.uuid4()) for _ in range(len(chunks))] # 各チャンクにユニークなIDを付与

            collection.add(
                documents=documents,
                metadatas=metadatas,
                ids=ids
            )
            processed_files_count += 1
            total_chunks_added += len(chunks)
            all_status_messages.append(f"PDFファイル '{file_name}' の処理が完了しました。{len(chunks)}個のチャンクがデータベースに登録されました。")

        except Exception as e:
            all_status_messages.append(f"PDFファイル '{os.path.basename(pdf_file.name)}' 処理中に予期せぬエラーが発生しました: {e}")
            continue # 次のファイルへ

    final_status_message = f"{processed_files_count}個のPDFファイルの処理が完了しました。合計{total_chunks_added}個のチャンクがデータベースに登録されました。質問を入力してください。\n\n" + "\n".join(all_status_messages)
    return final_status_message, gr.update(interactive=True) # 質問入力欄を有効化

def answer_question(question):
    """ChromaDBから関連情報を取得し、GPTモデルで質問に回答する"""
    if not question:
        return "質問を入力してください。","" # コンテキストも返すように変更

    if collection.count() == 0:
        return "PDFがまだアップロードされていないか、処理されていません。まずPDFをアップロードしてください。"

    try:
        # 1. 質問を埋め込みに変換し、ChromaDBを検索
        # ChromaDBのqueryメソッドは、埋め込みまたはテキストを受け取ります。
        # ここではテキストを直接渡して、ChromaDBのembedding_functionに処理させます。
        results = collection.query(
            query_texts=[question],
            n_results=8 # 上位5つの関連チャンクを取得
        )

        # 2. 検索結果からコンテキストを構築
        context_chunks = results['documents'][0] if results['documents'] else []
        if not context_chunks:
            return "関連する情報が見つかりませんでした。質問を明確にするか、別のPDFを試してください。"

        context = "\n\n".join(context_chunks)

        # 3. GPTモデルで回答を生成
        answer = get_gpt_response(question, context)
        return answer, context # 回答とコンテキストを両方返す
    except Exception as e:
        return f"質問応答中に予期せぬエラーが発生しました: {e}"

# --- Gradio UIの構築 ---
with gr.Blocks() as gradioUI:
    gr.Markdown(
        """
        # PDF Q&A with GPT-4o-mini and Vector Database
        PDFファイルをアップロードし、その内容について質問してください。
        **複数のPDFファイルを同時にアップロードできます。**
        "gpt-4o-mini"を使用しています。
        """
    )

    with gr.Row():
        with gr.Column():
            # file_count="multiple" を追加して複数ファイル対応にする
            pdf_input = gr.File(label="PDFドキュメントをアップロード", file_types=[".pdf"], file_count="multiple")
            upload_status = gr.Textbox(label="ステータス", interactive=False, value="PDFをアップロードしてください。", lines=5)

        with gr.Column():
            question_input = gr.Textbox(label="質問を入力", placeholder="このドキュメントは何について書かれていますか？", interactive=False)
            answer_output = gr.Markdown(label="回答")
            # 動作確認用コンポーネントを追加
            retrieved_context_output = gr.Textbox(label="取得されたコンテキスト", interactive=False, lines=10)

    # コンポーネント定義後にイベントリスナーを設定する
    # PDFアップロード時に処理を実行
    pdf_input.upload(
        upload_pdf_and_process,
        inputs=[pdf_input],
        outputs=[upload_status, question_input] # question_input が定義済みなのでエラーにならない
    )

    # 質問入力時に処理を実行
    question_input.submit(
        answer_question,
        inputs=[question_input],
        outputs=[answer_output, retrieved_context_output] # 2つの出力を受け取るように変更
    )

gradioUI.launch()