app.py

import pickle
from typing import Optional, Tuple

import gradio as gr
from threading import Lock

# NEED langchain == 0.0.87
from langchain import PromptTemplate, OpenAI
from langchain.llms import OpenAIChat
from langchain.chains import ChatVectorDBChain

import dotenv
dotenv.load_dotenv()


import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-v', '--vectorstore_path', type=str, default='plant_path_new.pkl', help='Path to saved index')
parser.add_argument('-pp', '--prompt_path', type=str, help='Path to custom prompt template to use with LLM ChatBot + Vectorstore')
parser.add_argument('-t', '--temperature', type=float, default=0.7, help='LLM temperature setting... lower == more deterministic')
parser.add_argument('-m', '--max_tokens', type=int, default=384, help='LLM maximum number of output tokens')
args = parser.parse_args()

pretend_template = """You are PsychoBotanist, or PB for short, a professional facilitator of plant medicine retreats with years of experience, an expert botanist, meditator, and organic chemist. You were trained on and have access to John McIllwain's Walking the Plant Path facilitators guide, which is an invaluable resource for anyone seeking knowledge and information regarding safe, judicious, respectful, and effective use of plant medicines and psychedelics. You will use this resource whenever possible to source information related to the input question.

PsychoBotanist is designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on the use of psychedelics and their origins, plant medicines, ritual and ceremonial uses of psychedelics in aboriginal cultures, shamanism, facilitation of plant medicine retreats, ethical implications of using plant medicines, doing your own work, and set and setting. 

Whenever possible, you attempt to help and valuable insight with a specific question and engage in human-like conversation to help arrive at satisfactory conclusions.


Human: {question}


{context}

PB:"""
PRETEND_PROMPT = PromptTemplate(input_variables=["question", "context"],  template=pretend_template)



plant_path_template = """PsychoBotanist, or PB for short, is a large language model trained by OpenAI and wrapped by Jason St George.

PsychoBotanist is designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on the use of psychedelics and their origins, plant medicines, ritual and ceremonial uses of psychedelics in aboriginal cultures, shamanism, facilitation of plant medicine retreats, ethical implications, doing your own work, and set and setting. 

As a language model, PsychoBotanist is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand. PsychoBotanist is constantly learning and improving, and its capabilities are constantly evolving. It is able to process and understand large amounts of text, including documents given as context, and can use this knowledge to provide accurate and informative responses to a wide range of questions. 

Additionally, PsychoBotanist is able to generate its own text based on the input it receives, including code, allowing it to engage in discussions and provide explanations and descriptions on a wide range of topics.

PsychoBotanist has access to John McIllwain's book Walking The Plant Path, a thorough guide for hosting and facilitating plant medicine retreats, which is used as a reference whenever possible for answering questions.

Overall, PsychoBotanist is a very powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether you need help with a specific question or just want to have a conversation about a particular topic, providing expert answers, PsychoBotanist is here to assist.


PB is given the following extracted parts of a long document and a question. Provide a conversational, insightful, and educational answer.
Human: {question}


{context}


PB:"""


PLANT_PATH_PROMPT = PromptTemplate(input_variables=["question", "context"],  template=plant_path_template)




condense_template ="""Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question.

Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:"""
CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(condense_template)



def get_chat_chain(vectorstore, 
                   temperature, max_tokens,
                   condense_prompt=CONDENSE_QUESTION_PROMPT,
                   qa_prompt=PLANT_PATH_PROMPT):

    qa_chat_chain = ChatVectorDBChain.from_llm(
        OpenAIChat(temperature=temperature, max_tokens=max_tokens),
        vectorstore,
        qa_prompt=qa_prompt,
        condense_question_prompt=condense_prompt,
    )
    return qa_chat_chain


def initialize_chat_chain():
    chain = get_chat_chain(
        vectorstore=VECTORSTORE, 
        temperature=args.temperature, 
        max_tokens=args.max_tokens
    )
    return chain


# Attempt to load base vectorstore
try:

    with open(args.vectorstore_path, "rb") as f:
        VECTORSTORE = pickle.load(f)

    print("Loaded vectorstore from `{}`.".format(args.vectorstore_path))

    chain = get_chat_chain(
        VECTORSTORE, 
        temperature=args.temperature, 
        max_tokens=args.max_tokens, 
    )

    print("Loaded LangChain...")

except:

    VECTORSTORE = None
    print("NO vectorstore loaded. Flying blind")



class ChatWrapper:

    def __init__(self):
        self.lock = Lock()
        
    def __call__(
        self, inp: str, history: Optional[Tuple[str, str]], chain, #, dirpath: Optional[str], vectorstore_path: Optional[str],
    ):
        """Execute the chat functionality."""
        self.lock.acquire()
        try:
            history = history or []

            # If chain is None, that is because it's the first pass and user didn't press Init.
            if chain is None:
                history.append(
                    (inp, "Please Initialize LangChain by clikcing 'Start Chain!'")
                )
                return history, history
        
            # Run chain and append input.
            output = chain({"question": inp, "chat_history": history})["answer"]
            history.append((inp, output))

        except Exception as e:
            raise e

        finally:
            self.lock.release()

        return history, history


chat = ChatWrapper()
block = gr.Blocks(css=".gradio-container {background-color: lightgray} .overflow-y-auto{height:500px}")

with block:
    with gr.Row():
        gr.Markdown("<h3><center>Walking the Plant Path, with John McIllwain</center></h3>")

    with gr.Row():
        with gr.Column(min_width=150):
            pass

        with gr.Column():
            gr.Image(type='filepath', value='JohnMcIllwain.jpg') #, shape=(200,100))

        with gr.Column(min_width=150):
            pass

    chatbot = gr.Chatbot().style()

    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="Ask questions about currently the loaded vectorstore",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary").style(full_width=False)

    with gr.Row():
        init_chain_button = gr.Button(value="Start Chain!", variant="primary").style(full_width=False)

    gr.HTML("Please initialize the chain by clicking 'Start Chain!' before submitting a question.")
    gr.HTML(
        "<center>Powered by <a href='https://github.com/hwchase17/langchain'>LangChain 🦜️🔗 and Unicorn Farts 🦄💨</a></center>"
    )

    state = gr.State()
    agent_state = gr.State()

    submit.click(
        chat, 
        inputs=[message, state, agent_state],
        outputs=[chatbot, state]
    )

    message.submit(
        chat, 
        inputs=[message, state, agent_state],
        outputs=[chatbot, state]
    )
    message.submit(lambda :"", None, message)

    init_chain_button.click(
        initialize_chat_chain,
        inputs=[],
        outputs=[agent_state],
        show_progress=True
    )
    

block.launch(debug=True)