OpenAITools/CrinicalTrialTools.py

from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import ChatOpenAI
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
# 修正: pydantic v1の非推奨警告を解決
try:
    from pydantic import BaseModel, Field
except ImportError:
    from langchain_core.pydantic_v1 import BaseModel, Field

import pandas as pd
import time
import re
from typing import Optional

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
# 修正: HuggingFaceEmbeddingsの非推奨警告を解決
try:
    from langchain_community.embeddings import HuggingFaceEmbeddings
except ImportError:
    from langchain.embeddings import HuggingFaceEmbeddings

from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# GPT初期化
gpt = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

## Database関連のインポート
from langchain_community.utilities import SQLDatabase
from sqlalchemy import create_engine
from langchain.prompts import ChatPromptTemplate
from langchain.schema import SystemMessage
from langchain_core.prompts import MessagesPlaceholder

from OpenAITools.FetchTools import fetch_clinical_trials, fetch_clinical_trials_jp

# エラーハンドリング用のデコレータ
def retry_on_error(max_retries=3, delay=1):
    """エラー時のリトライデコレータ"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    if attempt < max_retries - 1:
                        print(f"エラーが発生しました (試行 {attempt + 1}/{max_retries}): {e}")
                        if "no healthy upstream" in str(e) or "InternalServerError" in str(e):
                            time.sleep(delay * 2)  # サーバーエラーの場合は長めに待機
                        else:
                            time.sleep(delay)
                        continue
                    else:
                        print(f"最大リトライ回数に達しました: {e}")
                        raise e
            return None
        return wrapper
    return decorator

## Cancer Name の抽出
class ExtractTumorName(BaseModel):
    """Extract tumor name from the user's question."""
    tumor_name: str = Field(description="Extracted tumor name from the question, or 'None' if no tumor found")

class TumorNameExtractor:
    def __init__(self, llm):
        self.llm = llm
        self.structured_llm_extractor = self.llm.with_structured_output(ExtractTumorName)

        self.system_prompt = """あなたは、ユーザーの質問に基づいて腫瘍名を英語で抽出するシステムです。
質問文に腫瘍の種類や名前が含まれている場合、それを英語で返してください。
質問文に腫瘍名がない場合は 'None' と返答してください。"""

        self.grade_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    @retry_on_error(max_retries=3, delay=2)
    def extract_tumor_name(self, question: str) -> str:
        """腫瘍名を抽出するメソッド"""
        try:
            tumor_extractor = self.grade_prompt | self.structured_llm_extractor
            result = tumor_extractor.invoke({"question": question})
            return result.tumor_name
        except Exception as e:
            print(f"腫瘍名抽出エラー: {e}")
            return "None"

### 質問変更システム
class ModifyQuestion(BaseModel):
    """Class for modifying a question by inserting NCTID."""
    modified_question: str = Field(description="The modified question with the inserted NCTID.")

class QuestionModifier:
    def __init__(self, llm):
        self.llm = llm
        self.structured_llm_modifier = self.llm.with_structured_output(ModifyQuestion)

        # 修正: 中括弧をエスケープ
        self.system_prompt = """あなたは、ユーザーの質問に対して適切なNCTIDを挿入して質問を変更するシステムです。
質問文にNCTIDを挿入し、形式に基づいて新しい質問を生成してください。
例えば16歳男性の神経膠腫の患者さんが参加できる臨床治験を教えて下さいという質問に対しては
16歳男性の神経膠腫の患者さんは{{nct_id}}に参加できますか？と変更して下さい
NCTIDは {{nct_id}} を使用してください。"""

        self.modify_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    @retry_on_error(max_retries=3, delay=2)
    def modify_question(self, question: str, nct_id: str) -> str:
        """質問を変更するメソッド"""
        try:
            question_modifier = self.modify_prompt | self.structured_llm_modifier
            result = question_modifier.invoke({"question": question, "nct_id": nct_id})
            return result.modified_question
        except Exception as e:
            print(f"質問変更エラー: {e}")
            return f"{question} (NCTID: {nct_id})"

class QuestionModifierSecond:
    def __init__(self, llm):
        self.llm = llm
        self.structured_llm_modifier = self.llm.with_structured_output(ModifyQuestion)

        self.system_prompt = """あなたは、ユーザーの質問を変更するシステムです。
形式に基づいて新しい質問を生成してください。
例えば16歳男性の神経膠腫の患者さんが参加できる臨床治験を教えて下さいという質問に対しては
16歳男性の神経膠腫の患者さんはこの治験に参加できますか？と変更して下さい"""

        self.modify_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    @retry_on_error(max_retries=3, delay=2)
    def modify_question(self, question: str) -> str:
        """質問を変更するメソッド"""
        try:
            question_modifier = self.modify_prompt | self.structured_llm_modifier
            result = question_modifier.invoke({"question": question})
            return result.modified_question
        except Exception as e:
            print(f"質問変更エラー: {e}")
            return question

class QuestionModifierEnglish:
    def __init__(self, llm):
        self.llm = llm
        self.structured_llm_modifier = self.llm.with_structured_output(ModifyQuestion)

        self.system_prompt = """あなたは、ユーザーの質問を変更し英語に翻訳するシステムです。
形式に基づいて新しい質問を生成してください。
例えば16歳男性の神経膠腫の患者さんが参加できる臨床治験を教えて下さいという質問に対しては
Can a 16 year old male patient with glioma participate in this clinical trial?と変更して下さい"""

        self.modify_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    @retry_on_error(max_retries=3, delay=2)
    def modify_question(self, question: str) -> str:
        """質問を変更するメソッド"""
        try:
            question_modifier = self.modify_prompt | self.structured_llm_modifier
            result = question_modifier.invoke({"question": question})
            return result.modified_question
        except Exception as e:
            print(f"英語質問変更エラー: {e}")
            return question

### Make criteria check Agent
class ClinicalTrialAgent:
    def __init__(self, llm, db):
        self.llm = llm
        self.db = db

        self.system_prompt = """
        あなたは患者さんに適した治験を探すエージェントです。
        データベースのEligibility Criteriaをチェックして患者さんがその治験を受けることが可能かどうか答えて下さい
        """

        self.prompt = ChatPromptTemplate.from_messages(
            [("system", self.system_prompt), 
             ("human", "{input}"), 
             MessagesPlaceholder("agent_scratchpad")]
        )

        self.agent_executor = self.create_sql_agent(self.llm, self.db, self.prompt)

    def create_sql_agent(self, llm, db, prompt):
        """SQLエージェントを作成するメソッド"""
        agent_executor = create_sql_agent(
            llm, 
            db=db, 
            prompt=prompt, 
            agent_type="tool-calling", 
            verbose=True
        )
        return agent_executor

    @retry_on_error(max_retries=3, delay=2)
    def get_agent_judgment(self, modify_question: str) -> str:
        """Modifyされた質問を元に、患者さんが治験に参加可能かどうかのエージェント判断を取得"""
        try:
            result = self.agent_executor.invoke({"input": modify_question})
            return result
        except Exception as e:
            print(f"エージェント判断エラー: {e}")
            return f"エラー: {str(e)}"

class SimpleClinicalTrialAgent:
    def __init__(self, llm):
        self.llm = llm

    @retry_on_error(max_retries=3, delay=2)
    def evaluate_eligibility(self, TargetCriteria: str, question: str) -> str:
        """臨床試験の参加適格性を評価するメソッド"""
        try:
            # 修正: プロンプト内の中括弧を適切にエスケープ
            prompt_template = """
You are an agent looking for a suitable clinical trial for a patient.
Please answer whether the patient is eligible for this clinical trial based on the following criteria. If you do not know the answer, say you do not know. Your answer should be brief, no more than 3 sentences.

Question: {{question}} 

Criteria:
{criteria}""".format(criteria=TargetCriteria)

            criteria_prompt = ChatPromptTemplate.from_messages(
                [
                    (
                        "human",
                        prompt_template
                    )
                ]
            )

            rag_chain = (
                {"question": RunnablePassthrough()}
                | criteria_prompt
                | self.llm
                | StrOutputParser()
            )

            response = rag_chain.invoke(question)
            return response
            
        except Exception as e:
            print(f"適格性評価エラー: {e}")
            return f"評価エラー: {str(e)}"

### output 評価システム
class TrialEligibilityGrader(BaseModel):
    """3値評価: yes, no, unclear"""
    score: str = Field(
        description="The eligibility of the patient for the clinical trial based on the document. Options are: 'yes', 'no', or 'unclear'."
    )

class GraderAgent:
    def __init__(self, llm):
        self.llm = llm
        self.structured_llm_grader = self.llm.with_structured_output(TrialEligibilityGrader)

        self.system_prompt = """
        あなたは治験に参加する患者の適合性を評価するGraderです。
        以下のドキュメントを読み、患者が治験に参加可能かどうかを判断してください。
        'yes'（参加可能）、'no'（参加不可能）、'unclear'（判断できない）の3値で答えてください。
        """

        self.grade_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                (
                    "human",
                    "取得したドキュメント: \n\n {document} ",
                ),
            ]
        )

    @retry_on_error(max_retries=3, delay=2)
    def evaluate_eligibility(self, AgentJudgment_output: str) -> str:
        """AgentJudgment['output']を基に患者が治験に参加可能かどうかを評価し、スコア（AgentGrade）を返す"""
        try:
            GraderAgent = self.grade_prompt | self.structured_llm_grader
            result = GraderAgent.invoke({"document": AgentJudgment_output})
            return result.score
        except Exception as e:
            print(f"グレード評価エラー: {e}")
            return "unclear"

class LLMTranslator:
    def __init__(self, llm):
        self.llm = llm
        self.structured_llm_modifier = self.llm.with_structured_output(ModifyQuestion)

        self.system_prompt = """あなたは、優秀な翻訳者です。
日本語を英語に翻訳して下さい。"""

        self.system_prompt2 = """あなたは、優秀な翻訳者です。
日本語を英語に以下のフォーマットに従って翻訳して下さい。
MainQuestion:
Known gene mutation:
Measurable tumour:
Biopsiable tumour:"""

        self.modify_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

        self.modify_prompt2 = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt2),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    def is_english(self, text: str) -> bool:
        """簡易的にテキストが英語かどうかを判定する関数"""
        return bool(re.match(r'^[A-Za-z0-9\s.,?!]+$', text))

    @retry_on_error(max_retries=3, delay=2)
    def translate(self, question: str) -> str:
        """質問を翻訳するメソッド。英語の質問はそのまま返す。"""
        try:
            if self.is_english(question):
                return question
            
            question_modifier = self.modify_prompt | self.structured_llm_modifier
            result = question_modifier.invoke({"question": question})
            return result.modified_question
        except Exception as e:
            print(f"翻訳エラー: {e}")
            return question

    @retry_on_error(max_retries=3, delay=2)
    def translateQuestion(self, question: str) -> str:
        """フォーマット付きで質問を翻訳するメソッド"""
        try:
            question_modifier = self.modify_prompt2 | self.structured_llm_modifier
            result = question_modifier.invoke({"question": question})
            return result.modified_question
        except Exception as e:
            print(f"フォーマット翻訳エラー: {e}")
            return question

def generate_ex_question(age, sex, tumor_type, GeneMutation, Meseable, Biopsiable):
    """日本語での質問文を生成"""
    try:
        # GeneMutationが空の場合はUnknownに設定
        gene_mutation_text = GeneMutation if GeneMutation else "Unknown"
        
        # MeseableとBiopsiableの値をYes, No, Unknownに変換
        meseable_text = (
            "Yes" if Meseable == "有り" else "No" if Meseable == "無し" else "Unknown"
        )
        biopsiable_text = (
            "Yes" if Biopsiable == "有り" else "No" if Biopsiable == "無し" else "Unknown"
        )
        
        # 質問文の生成
        ex_question = f"""{age}歳{sex}の{tumor_type}患者さんはこの治験に参加することができますか？
判明している遺伝子変異: {gene_mutation_text}
Meseable tumor: {meseable_text}
Biopsiable tumor: {biopsiable_text}
です。"""
        return ex_question
    except Exception as e:
        print(f"日本語質問生成エラー: {e}")
        return f"{age}歳{sex}の{tumor_type}患者さんの治験参加について"

def generate_ex_question_English(age, sex, tumor_type, GeneMutation, Meseable, Biopsiable):
    """英語での質問文を生成"""
    try:
        # GeneMutationが空の場合は"Unknown"に設定
        gene_mutation_text = GeneMutation if GeneMutation else "Unknown"
        
        # sexの値を male または female に変換
        sex_text = "male" if sex == "男性" else "female" if sex == "女性" else "Unknown"
        
        # MeseableとBiopsiableの値を "Yes", "No", "Unknown" に変換
        meseable_text = (
            "Yes" if Meseable == "有り" else "No" if Meseable == "無し" else "Unknown"
        )
        biopsiable_text = (
            "Yes" if Biopsiable == "有り" else "No" if Biopsiable == "無し" else "Unknown"
        )
        
        # 英語での質問文を生成
        ex_question = f"""Can a {age}-year-old {sex_text} patient with {tumor_type} participate in this clinical trial?
Known gene mutation: {gene_mutation_text}
Measurable tumor: {meseable_text}
Biopsiable tumor: {biopsiable_text}"""
        return ex_question
    except Exception as e:
        print(f"英語質問生成エラー: {e}")
        return f"Can a {age}-year-old patient with {tumor_type} participate in this clinical trial?"

# テスト関数
def test_clinical_trial_tools():
    """臨床試験ツールのテスト関数"""
    try:
        from langchain_groq import ChatGroq
        
        # Groqクライアントの初期化
        groq = ChatGroq(model_name="llama3-70b-8192", temperature=0)
        
        # 各エージェントの初期化テスト
        translator = LLMTranslator(groq)
        criteria_agent = SimpleClinicalTrialAgent(groq)
        grader_agent = GraderAgent(groq)
        
        print("✅ 全てのエージェントが正常に初期化されました")
        
        # サンプル質問の生成テスト
        sample_question = generate_ex_question_English(
            age="45",
            sex="女性",
            tumor_type="breast cancer",
            GeneMutation="HER2",
            Meseable="有り",
            Biopsiable="有り"
        )
        
        print(f"✅ サンプル質問生成成功: {sample_question}")
        return True
        
    except Exception as e:
        print(f"❌ テスト中にエラーが発生しました: {e}")
        return False

if __name__ == "__main__":
    print("ClinicalTrialTools のテストを開始します...")
    success = test_clinical_trial_tools()
    if success:
        print("✅ テストが正常に完了しました。")
    else:
        print("❌ テストでエラーが発生しました。")