import os
from llama_index.core.query_engine import NLSQLTableQueryEngine
import pickle
import pandas as pd
import sqlalchemy as sa
#from llama_index.llms.replicate import Replicate
from llama_index.llms.groq import Groq
from llama_index.llms.llama_cpp.llama_utils import messages_to_prompt, completion_to_prompt
from llama_index.core import SQLDatabase
import requests
import re
from llama_index.llms.openai import OpenAI
from OpenAITools.FetchTools import fetch_clinical_trials, fetch_clinical_trials_jp
import gradio as gr


def custom_completion_to_prompt(completion:str) ->str:
    return completion_to_prompt(
        completion, system_prompt=(
            "You are a sophisticated AI model designed to assist users in finding relevant clinical trials efficiently.  **Accurate Retrieval**: Utilize the RAG mechanism to retrieve the most relevant and current clinical trial information from reliable databases. Ensure the information is accurate and cite the sources when necessary. Make decisions based solely on the information provided by the retriever."
        ),
    )

def RetriviralClinicalTrrial_multi(message,history,cancer_name):
    df = fetch_clinical_trials(cancer_name)
    engine = sa.create_engine("sqlite:///:memory:")
    sql_database = df.to_sql(name='clinical_study', con=engine, if_exists='replace', index=False)
    database = SQLDatabase(engine)
    # The replicate endpoint
    LLAMA3_8B = "Llama3-8b-8192"
    LLAMA3_70b = "Llama3-70b-8192"
    Mixtral  = "mixtral-8x7b-32768" 

    '''llm = OpenAI(
    model='gpt-3.5-turbo',
    temperature=0.01,
    context_window=4096,
    completion_to_prompt=custom_completion_to_prompt,
    messages_to_prompt=messages_to_prompt,)'''

    llm = Groq(
        model=LLAMA3_70b,
        temperature=0.01,
        context_window=4096,
        completion_to_prompt=custom_completion_to_prompt,
        messages_to_prompt=messages_to_prompt,
    )
    query_engine = NLSQLTableQueryEngine(
    sql_database = database,
    tables=["clinical_study"],llm=llm)
    response = query_engine.query(message)
    return response.response

with gr.Blocks() as demo:
    cancer_name = gr.Textbox("lung cancer", label="cancer name")

    gr.ChatInterface(
        RetriviralClinicalTrrial_multi, additional_inputs=[cancer_name],
        title="Assistant for Clinical traial Search",
        description='リアルタイムで日本で行われているClinical trialの情報を入手して治験を探すお手伝いをします。対象となるがんの名前を上に入力して、まずNCTIDを聞いてください。その後にそれぞれのNCTIDのを聴いてください。その後興味のある内容　Title,Cancer,Summary,Japanes Locations,Eligibility Criteriaなどを聞いてみてください（スペルの打ち間違いに注意）例えば「NCT04270591のTitleを教えて』のように尋ねてください',
        theme="soft",
        examples=[["NCTIDを全て教えて"], ["のTitleを教えて"], ["の対象Cancerを教えて"], ["のSummaryを教えて"], ["のJapanes Locationsを教えて"], ["のEligibility Criteriaを教えて"], ["のPrimary Completion Dateを教えて"]]
    )

if __name__ == "__main__":
    demo.launch()