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 # --- Ollamaクライアントの初期化 --- client_ollama = ollama.Client() OLLAMA_MODEL_NAME = "llama3.2" # OLLAMA_MODEL_NAME = "llama3:8b-instruct-q4_0" # --- 埋め込みモデルの初期化 --- # 重複定義を削除し、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ファイルからテキストを抽出する""" try: reader = PdfReader(pdf_file_path) text = "" for page in reader.pages: page_text = page.extract_text() if page_text: # ページが空でなければ追加 text += page_text + "\n" if not text.strip(): # 全ページからテキストが抽出できなかった場合 return "PDFからテキストを抽出できませんでした。画像ベースのPDFかもしれません。" return text except Exception as e: return f"PDFの読み込み中にエラーが発生しました: {e}" def get_ollama_response(query, context, source_code_to_check): """Ollamaモデルを使用して質問に回答する""" messages = [ {"role": "system", "content": "あなたは提供されたコンテキスト（ソースコードチェックリスト）とレビュー対象のソースコードに基づいて、ソースコードをチェックし、その結果を返す有益なアシスタントです。チェックリストの項目ごとにソースコードを評価し、具体的な指摘と改善案を提示してください。コンテキストに情報がない場合は、「提供された情報からは回答できません。」と答えてください。"}, {"role": "user", "content": f"ソースコードチェックリスト:\n{context}\n\nレビュー対象のソースコード:\n```\n{source_code_to_check}\n```\n\n質問: {query}\n\nチェック結果:"} ] try: print(f"Ollamaモデル '{OLLAMA_MODEL_NAME}' にリクエストを送信中...") response = client_ollama.chat( model=OLLAMA_MODEL_NAME, messages=messages, options={ "temperature": 0.5, "num_predict": 2000 # 回答の最大トークン数 } ) print(f"Ollamaからの生応答: {response}") # デバッグ用にOllamaの生レスポンスを出力 if 'message' in response and 'content' in response['message']: return response['message']['content'].strip() else: print(f"Ollamaからの応答形式が不正です: {response}") return "Ollamaモデルからの応答形式が不正です。詳細をコンソールログで確認してください。" except ollama.ResponseError as e: # Ollama固有のエラーを捕捉 print(f"Ollama APIエラーが発生しました: {e}") return f"Ollamaモデルの呼び出し中にAPIエラーが発生しました: {e}\nOllamaサーバーが起動しているか、モデル '{OLLAMA_MODEL_NAME}' がインストールされているか確認してください。" except Exception as e: # その他の予期せぬエラー print(f"Ollamaモデルの呼び出し中に予期せぬエラーが発生しました: {e}") return f"Ollamaモデルの呼び出し中に予期せぬエラーが発生しました: {e}" def upload_pdf_and_process(pdf_files): """複数のPDFファイルをアップロードし、テキストを抽出し、ChromaDBに登録する""" if not pdf_files: 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}") # 1. PDFからテキストを抽出 raw_text = extract_text_from_pdf(pdf_path) if "エラー" in raw_text or "抽出できませんでした" in raw_text: all_status_messages.append(raw_text) print(f"Error extracting text from {file_name}: {raw_text}") continue # 次のファイルへ # 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): """ChromaDBから関連情報を取得し、Ollamaモデルで質問に回答する""" if not question and not source_code: # 質問またはソースコードのいずれかがあればOK 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 # 上位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_ollama_response(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ファイルを同時にアップロードできます。** ローカルのOllama ({OLLAMA_MODEL_NAME}) を使用しています。 """ ) 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(): source_code_input = gr.Code( label="レビュー対象のソースコード (ここにソースコードを貼り付けてください)", value="", language="python", interactive=True, lines=15 ) question_input = gr.Textbox(label="レビュー指示（例: セキュリティの観点からレビュー）", placeholder="特定の観点からのレビュー指示を入力してください（任意）。", interactive=False) 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], outputs=[answer_output, retrieved_context_output] ) gradioUI.launch(server_name="0.0.0.0", server_port=7860)