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
from langchain_core.pydantic_v1 import BaseModel, Field
import pandas as pd
from pydantic import BaseModel, Field

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser



gpt = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
#agent_gpt_executor = create_sql_agent(gpt, db=db, agent_type="tool-calling", verbose=True)

## make 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
#agent_groq_executor = create_sql_agent(llm, db=db, agent_type="tool-calling", verbose=True)

from OpenAITools.FetchTools import fetch_clinical_trials, fetch_clinical_trials_jp



## 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

        # LLMの出力を構造化するための設定
        self.structured_llm_extractor = self.llm.with_structured_output(ExtractTumorName)

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

        # プロンプトテンプレート
        self.grade_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    def extract_tumor_name(self, question: str) -> str:
        """
        腫瘍名を抽出するメソッド。
        :param question: 質問文
        :return: 抽出された腫瘍名
        """
        # 質問から腫瘍名を抽出する処理
        tumor_extractor = self.grade_prompt | self.structured_llm_extractor
        result = tumor_extractor.invoke({"question": question})
        return result.tumor_name
    
### 質問変更システム

# ModifyQuestionの出力形式を定義
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

        # LLMの出力を構造化するための設定
        self.structured_llm_modifier = self.llm.with_structured_output(ModifyQuestion)

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

        # プロンプトテンプレート
        self.modify_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    def modify_question(self, question: str, nct_id: str) -> str:
        """
        質問を変更するメソッド。
        :param question: 質問文
        :param nct_id: NCTID
        :return: NCTIDを挿入した新しい質問
        """
        # 質問を変更するプロセス
        question_modifier = self.modify_prompt | self.structured_llm_modifier
        result = question_modifier.invoke({"question": question, "nct_id": nct_id})
        modify_question = result.modified_question
        return modify_question

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

        # LLMの出力を構造化するための設定
        self.structured_llm_modifier = self.llm.with_structured_output(ModifyQuestion)

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

        # プロンプトテンプレート
        self.modify_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    def modify_question(self, question: str) -> str:
        """
        質問を変更するメソッド。
        :param question: 質問文
        :param nct_id: NCTID
        :return: NCTIDを挿入した新しい質問
        """
        # 質問を変更するプロセス
        question_modifier = self.modify_prompt | self.structured_llm_modifier
        result = question_modifier.invoke({"question": question})
        modify_question = result.modified_question
        return modify_question
    
class QuestionModifierEnglish:
    def __init__(self, llm):
        self.llm = llm

        # LLMの出力を構造化するための設定
        self.structured_llm_modifier = self.llm.with_structured_output(ModifyQuestion)

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

        # プロンプトテンプレート
        self.modify_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                ("human", "ユーザーの質問: {question}"),
            ]
        )

    def modify_question(self, question: str) -> str:
        """
        質問を変更するメソッド。
        :param question: 質問文
        :param nct_id: NCTID
        :return: NCTIDを挿入した新しい質問
        """
        # 質問を変更するプロセス
        question_modifier = self.modify_prompt | self.structured_llm_modifier
        result = question_modifier.invoke({"question": question})
        modify_question = result.modified_question
        return modify_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")]
        )

        # SQL Agentの設定
        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

    def get_agent_judgment(self, modify_question: str) -> str:
        """
        Modifyされた質問を元に、患者さんが治験に参加可能かどうかのエージェント判断を取得。
        :param modify_question: NCTIDが挿入された質問
        :return: エージェントの判断 (AgentJudgment)
        """
        # LLMに質問を投げて、判断を得る
        result = self.agent_executor.invoke({"input": modify_question})
        return result
    

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

    def evaluate_eligibility(self, TargetCriteria: str, question: str) -> str:
        """
        臨床試験の参加適格性を評価するメソッド。
        :param TargetCriteria: 試験基準 (Inclusion/Exclusion criteria)
        :param question: 患者の条件に関する質問
        :return: 臨床試験に参加可能かどうかのLLMからの応答
        """

        # プロンプトの定義
        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:
        """ + TargetCriteria

        # プロンプトテンプレートの作成
        criteria_prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "human",
                    prompt_template
                )
            ]
        )

        # RAGチェーンの作成
        rag_chain = (
            {"question": RunnablePassthrough()}
            | criteria_prompt
            | self.llm
            | StrOutputParser()
        )

        # 質問をチェーンに渡して、応答を得る
        response = rag_chain.invoke(question)
        return response
    

### 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

        # LLMの出力を構造化するための設定
        self.structured_llm_grader = self.llm.with_structured_output(TrialEligibilityGrader)

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

        # 評価のためのプロンプトを作成
        self.grade_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", self.system_prompt),
                (
                    "human",
                    "取得したドキュメント: \n\n {document} ",
                ),
            ]
        )

    def evaluate_eligibility(self, AgentJudgment_output: str) -> str:
        """
        AgentJudgment['output']を基に患者が治験に参加可能かどうかを評価し、スコア（AgentGrade）を返す。
        :param AgentJudgment_output: エージェント判断の 'output' の値
        :return: 評価されたスコア (AgentGrade)
        """
        GraderAgent = self.grade_prompt | self.structured_llm_grader
        result = GraderAgent.invoke({"document": AgentJudgment_output})
        AgentGrade = result.score
        return AgentGrade

import re

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

        self.system_prompt = """あなたは、優秀な翻訳者です。\n
                                日本語を英語に翻訳して下さい。\n
                             """
        self.system_prompt2 = """あなたは、優秀な翻訳者です。\n
                                日本語を英語に以下のフォーマットに従って翻訳して下さい。\n
                                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:
        """
        簡易的にテキストが英語かどうかを判定する関数
        :param text: 判定するテキスト
        :return: 英語の場合True、日本語の場合False
        """
        # 英語のアルファベットが多く含まれているかを確認
        return bool(re.match(r'^[A-Za-z0-9\s.,?!]+$', text))

    def translate(self, question: str) -> str:
        """
        質問を翻訳するメソッド。英語の質問はそのまま返す。
        :param question: 質問文
        :return: 翻訳済みの質問文、または元の質問文（英語の場合）
        """
        # 質問が英語の場合、そのまま返す
        if self.is_english(question):
            return question
        
        # 日本語の質問は翻訳プロセスにかける
        question_modifier = self.modify_prompt | self.structured_llm_modifier
        result = question_modifier.invoke({"question": question})
        modify_question = result.modified_question
        return modify_question
    
    def translateQuestion(self, question: str) -> str:
        """
        フォーマット付きで質問を翻訳するメソッド。
        :param question: 質問文
        :return: フォーマットに従った翻訳済みの質問
        """
        question_modifier = self.modify_prompt2 | self.structured_llm_modifier
        result = question_modifier.invoke({"question": question})
        modify_question = result.modified_question
        return modify_question

def generate_ex_question(age, sex, tumor_type, GeneMutation, Meseable, Biopsiable):
    # 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

def generate_ex_question_English(age, sex, tumor_type, GeneMutation, Meseable, Biopsiable):
    # 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