from paperqa import Docs
import pickle

from ring import LLM_GPT35
from ring import LLM_GPT4
from ring import LLM_LLAMA2
from ring import get_llm

my_docs = [
    "corpus/MirzakhaliliEtAl2018.pdf",
    "corpus/NicollettiEtAl2019.pdf",
    "corpus/ATributetoSydneyBrenner.pdf",
]

# my_docs = ['corpus/ATributetoSydneyBrenner.pdf']

llm_to_use = LLM_GPT4
llm = get_llm(llm_to_use, 0.1)

CORPUS_PICKLE = "corpus.pkl"


def generate_pickle():
    docs = Docs(llm="langchain", client=llm)

    print("Initialised Docs with LLM: %s" % llm_to_use)

    for d in my_docs:
        print("Adding %s" % d)
        docs.add(d)

    print("Saving pickle of docs to: %s" % CORPUS_PICKLE)

    f = open(CORPUS_PICKLE, "wb")
    pickle.dump(docs, f)
    f.close()


def run_query(query, llm_ver):

    f = open(CORPUS_PICKLE, "rb")

    docs = pickle.load(f)

    llm = get_llm(llm_ver, 0)

    docs.set_client(llm)
    answer = docs.query(query)

    return answer


if __name__ == "__main__":

    # generate_pickle()

    print("Running queries")

    queries = [
        "Where did John Sulston do his postdoc?",
        "What types of potassium channels are expressed in C. elegans neurons?",
    ]

    for query in queries:
        print("-------------------")
        print("Q: %s" % query)
        # answer = docs.query(query)
        answer = run_query(query, LLM_GPT4)

        print("A: %s" % answer)