app_multillm_OK.py

import os
import gradio as gr
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.utils import embedding_functions
import ollama # Ollamaライブラリをインポート
from langchain_text_splitters import RecursiveCharacterTextSplitter
import uuid # For generating unique IDs for chunks
from dotenv import load_dotenv
# from pdfminer.high_level import extract_text as pdfminer_extract_text


# LLMクライアントのインポート
from openai import OpenAI
import anthropic
import google.generativeai as genai

import sys

print(f"Python executable: {sys.executable}")
print(f"Python version: {sys.version}")
print(f"Python version info: {sys.version_info}")
print(f"--------------------------")

# .envファイルから環境変数を読み込む
load_dotenv()

# --- APIキーの取得 ---
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# --- Ollamaクライアントの初期化 ---
client_ollama = ollama.Client()

OLLAMA_MODEL_NAME = "llama3.2"
# OLLAMA_MODEL_NAME = "llama3:8b-instruct-q4_0"


client_openai = None
OPENAI_MODEL_NAME = "gpt-4o-mini"
if OPENAI_API_KEY:
    try:
        client_openai = OpenAI(api_key=OPENAI_API_KEY)
        print(f"OpenAIクライアントを初期化しました (モデル: {OPENAI_MODEL_NAME})。")
    except Exception as e:
        print(f"OpenAIクライアントの初期化に失敗しました: {e}")
        client_openai = None
else:
    print("OPENAI_API_KEYが設定されていません。OpenAIモデルは利用できません。")

client_anthropic = None
ANTHROPIC_MODEL_NAME = "claude-3-haiku-20240307"
if ANTHROPIC_API_KEY:
    try:
        client_anthropic = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)
        print(f"Anthropicクライアントを初期化しました (モデル: {ANTHROPIC_MODEL_NAME})。")
    except Exception as e:
        print(f"Anthropicクライアントの初期化に失敗しました: {e}")
        client_anthropic = None
else:
    print("ANTHROPIC_API_KEYが設定されていません。Anthropicモデルは利用できません。")

client_gemini = None
GOOGLE_MODEL_NAME = "gemini-2.5-flash"
if GOOGLE_API_KEY:
    try:
        genai.configure(api_key=GOOGLE_API_KEY)
        client_gemini = genai.GenerativeModel(GOOGLE_MODEL_NAME)
        print(f"Google Geminiクライアントを初期化しました (モデル: {GOOGLE_MODEL_NAME})。")
    except Exception as e:
        print(f"Google Geminiクライアントの初期化に失敗しました: {e}")
        client_gemini = None
else:
    print("GOOGLE_API_KEYが設定されていません。Google Geminiモデルは利用できません。")

# --- 埋め込みモデルの初期化 ---
# 重複定義を削除し、1回のみ初期化
embedding_model = SentenceTransformer('pkshatech/GLuCoSE-base-ja') # 日本語対応の埋め込みモデル

# --- ChromaDBのカスタム埋め込み関数 ---
# 重複定義を削除し、1回のみ定義
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クライアントとコレクションの初期化 ---
# インメモリモードで動作させ、アプリケーション起動時にコレクションをリセットします。
# グローバル変数としてクライアントを保持
client = chromadb.Client()
collection_name = "pdf_documents_collection"

# アプリケーション起動時にコレクションが存在すれば削除し、新しく作成する
# （インメモリDBはセッションごとにリセットされるため、これは初回起動時のみ意味を持つ）
try:
    client.delete_collection(name=collection_name)
    print(f"既存のChromaDBコレクション '{collection_name}' を削除しました。")
except Exception as e:
    # コレクションが存在しない場合はエラーになるので無視。デバッグ用にメッセージは出力。
    print(f"ChromaDBコレクション '{collection_name}' の削除に失敗しました (存在しないか、その他のエラー): {e}")
    pass

collection = client.get_or_create_collection(name=collection_name, embedding_function=sbert_ef)
print(f"ChromaDBコレクション '{collection_name}' を初期化しました。")

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ファイルからテキストを抽出する"""
    print(f"Attempting to extract text from: {pdf_file_path}")
    try:
        reader = PdfReader(pdf_file_path)
        text = ""
        if not reader.pages:
            print(f"  PDF '{os.path.basename(pdf_file_path)}' contains no pages.")
            return "ERROR: PDFにページが含まれていません。" # プレフィックスを追加

        for i, page in enumerate(reader.pages):
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
                # print(f"  Page {i+1} extracted text (first 100 chars): {page_text[:100].replace('\n', ' ')}...")
                cleaned_page_text = page_text[:100].replace('\n', ' ')
                print(f"  Page {i+1} extracted text (first 100 chars): {cleaned_page_text}...")
            else:
                print(f"  Page {i+1} extracted no text.")

        if not text.strip():
            print("  No text extracted from any page.")
            return "ERROR: PDFからテキストを抽出できませんでした。画像ベースのPDFかもしれません。" # プレフィックスを追加

        print(f"  Total text extracted (length: {len(text)}).")
        return text
    except Exception as e:
        print(f"  Error during PDF reading: {e}")
        return f"ERROR: PDFの読み込み中にエラーが発生しました: {e}" # プレフィックスを追加

def get_llm_response(selected_llm, query, context, source_code_to_check):
    """選択されたLLMを使用して質問に回答する"""
    system_prompt = "あなたは提供されたコンテキスト（ソースコードチェックリスト）とレビュー対象のソースコードに基づいて、ソースコードをチェックし、その結果を返す有益なアシスタントです。チェックリストの項目ごとにソースコードを評価し、具体的な指摘と改善案を提示してください。コンテキストに情報がない場合は、「提供された情報からは回答できません。」と答えてください。"
    user_content = f"ソースコードチェックリスト:\n{context}\n\nレビュー対象のソースコード:\n```\n{source_code_to_check}\n```\n\nレビュー指示: {query}\n\nチェック結果:"

    try:
        if selected_llm == "Ollama":
            if not client_ollama:
                return "Ollamaクライアントが初期化されていません。"
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_content}
            ]
            print(f"Ollamaモデル '{OLLAMA_MODEL_NAME}' にリクエストを送信中...")
            response = client_ollama.chat(
                model=OLLAMA_MODEL_NAME,
                messages=messages,
                options={
                    "temperature": 0.7,
                    "num_predict": 2000
                }
            )
            if 'message' in response and 'content' in response['message']:
                return response['message']['content'].strip()
            else:
                return f"Ollamaからの応答形式が不正です: {response}"

        elif selected_llm == "GPT":
            if not client_openai:
                return "OpenAI APIキーが設定されていないため、GPTモデルは利用できません。"
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_content}
            ]
            print(f"GPTモデル '{OPENAI_MODEL_NAME}' にリクエストを送信中...")
            response = client_openai.chat.completions.create(
                model=OPENAI_MODEL_NAME,
                messages=messages,
                temperature=0.5,
                max_tokens=2000
            )
            return response.choices[0].message.content.strip()

        elif selected_llm == "Anthropic":
            if not client_anthropic:
                return "Anthropic APIキーが設定されていないため、Anthropicモデルは利用できません。"
            messages = [
                {"role": "user", "content": user_content}
            ]
            print(f"Anthropicモデル '{ANTHROPIC_MODEL_NAME}' にリクエストを送信中...")
            response = client_anthropic.messages.create(
                model=ANTHROPIC_MODEL_NAME,
                max_tokens=2000,
                temperature=0.5,
                system=system_prompt, # Anthropicはsystemプロンプトを直接引数で渡す
                messages=messages
            )
            return response.content[0].text.strip()

        elif selected_llm == "Google Gemini":
            if not client_gemini:
                return "Google APIキーが設定されていないため、Geminiモデルは利用できません。"
            # Geminiのsystem instructionはまだベータ版で、messagesと併用できない場合があるため、
            # system_promptをuser_contentの先頭に結合する形式にする。
            # --- システムプロンプトの調整 (後述の2.プロンプト調整も参照) ---
            system_prompt = "あなたは提供されたコンテキスト（ソースコードチェックリスト）とレビュー対象のソースコードに基づいて、ソースコードのレビューを行うアシスタントです。チェックリストの項目ごとにソースコードを評価し、潜在的な問題点や改善の機会を提案してください。コンテキストに情報がない場合は、「提供された情報からは回答できません。」と答えてください。"
            # --- ユーザープロンプトの調整 (後述の2.プロンプト調整も参照) ---
            user_content = f"ソースコードチェックリスト:\n{context}\n\nレビュー対象のソースコード:\n```\n{source_code_to_check}\n```\n\nレビュー指示: {query}\n\nチェック結果:"

            full_user_content = f"{system_prompt}\n\n{user_content}"
            messages = [
                {"role": "user", "parts": [full_user_content]}
            ]
            print(f"Google Geminiモデル '{GOOGLE_MODEL_NAME}' にリクエストを送信中...")
            try:
                response = client_gemini.generate_content(
                    messages,
                    generation_config=genai.types.GenerationConfig(
                        temperature=0.5, # まずは0.5で試す。必要なら0.7などに上げる
                        max_output_tokens=2000
                    )
                )

                # --- エラーハンドリングの強化 ---
                # response.text を呼び出す前に、応答の候補と終了理由を確認
                if response.candidates:
                    candidate = response.candidates[0]
                    # finish_reason が SAFETY (genai.types.HarmCategory.SAFETY) の場合、安全ポリシーによりブロックされた可能性が高い
                    if candidate.finish_reason == genai.types.HarmCategory.SAFETY:
                        safety_ratings = candidate.safety_ratings
                        safety_details = ", ".join([f"{sr.category.name}: {sr.probability.name}" for sr in safety_ratings])
                        print(f"Gemini response blocked due to safety policy. Details: {safety_details}")
                        return f"Google Geminiからの応答が安全ポリシーによりブロックされました。詳細: {safety_details}"
                    # 正常なコンテンツがあるか確認
                    elif candidate.content and candidate.content.parts:
                        return response.text.strip()
                    else:
                        # コンテンツがないが、finish_reasonがSAFETY以外の場合
                        print(f"Gemini response has no content parts. Finish reason: {candidate.finish_reason.name}")
                        return f"Google Geminiからの応答にコンテンツが含まれていません。終了理由: {candidate.finish_reason.name}"
                else:
                    # 候補自体がない場合
                    print(f"Gemini response has no candidates. Raw response: {response}")
                    return f"Google Geminiからの応答に候補がありませんでした。生の応答: {response}"

            except Exception as e:
                # generate_content 自体でエラーが発生した場合
                print(f"Google Gemini API呼び出し中にエラーが発生しました: {e}")
                return f"Google Gemini API呼び出し中にエラーが発生しました: {e}"

        else:
            return "無効なLLMが選択されました。"

    except Exception as e:
        print(f"LLM ({selected_llm}) の呼び出し中にエラーが発生しました: {e}")
        return f"LLM ({selected_llm}) の呼び出し中にエラーが発生しました: {e}"

def upload_pdf_and_process(pdf_files):
    """複数のPDFファイルをアップロードし、テキストを抽出し、ChromaDBに登録する"""
    if not pdf_files:
        print("No PDF files uploaded.")
        return "PDFファイルがアップロードされていません。", gr.update(interactive=False), gr.update(interactive=False)

    processed_files_count = 0
    total_chunks_added = 0
    all_status_messages = []

    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}' を処理中...")
            print(f"Processing PDF: {file_name} (Temporary Path: {pdf_path})")

            # 1. PDFからテキストを抽出
            raw_text = extract_text_from_pdf(pdf_path)

            # --- デバッグ用追加コード (前回のデバッグで追加したものは残しておくと良いでしょう) ---
            print(f"DEBUG: raw_text received from extract_text_from_pdf (length: {len(raw_text)})")
            # print(f"DEBUG: raw_text starts with: '{raw_text[:100].replace(newline_char, ' ')}'")
            print(f"DEBUG: 'エラー' in raw_text: {'エラー' in raw_text}")
            print(f"DEBUG: '抽出できませんでした' in raw_text: {'抽出できませんでした' in raw_text}")
            print(f"DEBUG: 'PDFにページが含まれていません' in raw_text: {'PDFにページが含まれていません' in raw_text}")
            # --- デバッグ情報ここまで ---

            # エラープレフィックスでチェックするように変更
            if raw_text.startswith("ERROR:"): # ここを変更
                all_status_messages.append(raw_text)
                print(f"Error during text extraction from {file_name}: {raw_text}") # ログメッセージも変更
                continue # 次のファイルへ

            # --- デバッグ用追加コード ---
            print(f"\n--- Raw text extracted from {file_name} (length: {len(raw_text)}, first 500 chars) ---")
            print(raw_text[:500])
            print(f"--- End of raw text from {file_name} ---\n")

            # 2. テキストをチャンクに分割
            chunks = text_splitter.split_text(raw_text)
            if not chunks:
                all_status_messages.append(f"'{file_name}' から有効なテキストチャンクを抽出できませんでした。")
                print(f"No valid chunks extracted from {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))]
            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)}個のチャンクがデータベースに登録されました。")
            print(f"Finished processing {file_name}. Added {len(chunks)} chunks.")

        except Exception as e:
            all_status_messages.append(f"PDFファイル '{os.path.basename(pdf_file.name)}' 処理中に予期せぬエラーが発生しました: {e}")
            print(f"Unexpected error during processing {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), gr.update(interactive=True)

def answer_question(question, source_code, selected_llm):
    """ChromaDBから関連情報を取得し、選択されたLLMで質問に回答する"""
    if not question and not source_code:
        return "質問またはレビュー対象のソースコードを入力してください。", ""

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

    try:
        print(f"Searching ChromaDB for question: {question}")
        results = collection.query(
            query_texts=[question],
            n_results=8
        )
        context_chunks = results['documents'][0] if results['documents'] else []
        if not context_chunks:
            print("No relevant context chunks found in ChromaDB.")
            return "関連する情報が見つかりませんでした。質問を明確にするか、別のPDFを試してください。", ""

        context = "\n\n".join(context_chunks)
        print(f"Retrieved context (first 500 chars):\n{context[:500]}...")

        answer = get_llm_response(selected_llm, question, context, source_code)
        return answer, context
    except Exception as e:
        print(f"質問応答中に予期せぬエラーが発生しました: {e}")
        return f"質問応答中に予期せぬエラーが発生しました: {e}", ""

# --- Gradio UIの構築 ---
with gr.Blocks() as gradioUI:
    gr.Markdown(
        f"""
        # PDF Q&A with Local LLM (Ollama: {OLLAMA_MODEL_NAME}) and Vector Database
        PDFファイルとしてソースコードチェックリストをアップロードし、レビューしたいソースコードを入力してください。
        **複数のPDFファイルを同時にアップロードできます。**
        利用するLLMを選択してください。
        """
        )
    with gr.Row():
        with gr.Column():
            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():
        # LLM選択コンポーネント
        llm_options = ["Ollama"]
        if client_openai:
            llm_options.append("GPT")
        if client_anthropic:
            llm_options.append("Anthropic")
        if client_gemini:
            llm_options.append("Google Gemini")

        llm_choice = gr.Radio(
            llm_options,
            label="使用するLLMを選択",
            value=llm_options[0] if llm_options else None, # 利用可能な最初のLLMをデフォルトにする
            interactive=True
        )

        source_code_input = gr.Code(
            label="レビュー対象のソースコード (ここにソースコードを貼り付けてください)",
            value="",
            language="python",
            interactive=False, # PDFアップロード後に有効化
            lines=15
        )
        question_input = gr.Textbox(label="レビュー指示（例: セキュリティの観点からレビュー）", placeholder="特定の観点からのレビュー指示を入力してください（任意）。", interactive=False) # PDFアップロード後に有効化

        review_button = gr.Button("レビュー開始")

        answer_output = gr.Markdown(label="レビュー結果")
        retrieved_context_output = gr.Textbox(label="取得されたチェックリスト項目", interactive=False, lines=10)

    pdf_input.upload(
        upload_pdf_and_process,
        inputs=[pdf_input],
        outputs=[upload_status, question_input, source_code_input]
    )

    review_button.click(
        answer_question,
        inputs=[question_input, source_code_input, llm_choice],
        outputs=[answer_output, retrieved_context_output]
    )

# gradioUI.launch(server_name="localhost", server_port=7860)
gradioUI.launch(server_name="0.0.0.0", server_port=7860)