agent.py

from dotenv import load_dotenv
#from dotenv import dotenv_values
import os
from tqdm.auto import tqdm
import pinecone
from langchain.embeddings import OpenAIEmbeddings
from pinecone_text.sparse import BM25Encoder
from langchain.retrievers import PineconeHybridSearchRetriever
from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent, Tool
from langchain.tools.base import BaseTool
from langchain.agents import AgentType
from langchain.agents.react.base import DocstoreExplorer
from langchain import LLMMathChain
from typing import Union
from langchain.memory import ConversationBufferWindowMemory
import random
from pydantic import Extra

from langchain.chains import LLMChain
from langchain import PromptTemplate




import promptlayer
from langchain.callbacks import PromptLayerCallbackHandler


load_dotenv()
#os.environ = dotenv_values(".env")

promptlayer.api_key = os.environ["PROPTLAYER_API_KEY"]

assistant_name = os.environ["ASSISTANT_NAME"]
topic = os.environ["TOPIC"]
course_name = os.environ["COURSE_NAME"]
institution = os.environ["INSTITUTION"]


class CalculatorTool(BaseTool):
    name = "CalculatorTool"
    
    description = """
    Useful for when you need to execute specific math calculations.
    This tool is only for math calculations and nothing else.
    Formulate the input as python code.
    """

    def _run(self, question: str):
        return exec(question)
    
    def _arun(self, question: str):
        raise NotImplementedError("This tool does not support async")


class GetrandomTool(BaseTool):
    name = "GetrandomTool"
    
    description = f"""
    Useful for when you need to get any randomly chosen piece of document regarding {topic} 
    from the study material. This is especially useful if the student wants tutoring,
    that is, he/she wants {assistant_name} to ask him/her questions about the study material.
    To use this tool, just call it with the constant text RANDOM.
    """
    class Config:
        extra = Extra.allow

    def _run(self, question: str):
        rand_id = str(random.randint(0,self.index_max))
        text = self.indexer.fetch([rand_id])["vectors"][rand_id]["metadata"]["context"]
        
        return text

    def _arun(self, question: str):
        raise NotImplementedError("This tool does not support async")


#Tutor tool?
#A chain with retrieval for answering, and a constant input summary of the
#tutoring flow so far

class TutoringTool(BaseTool):
    name = "TutoringTool"

    description = """This tool is capable of generating tutoring questions. 
    It has to be called with a summary of the previous tutoring discussion steps,
    or in case of a new tutoring session, with a randomly chosen piece of material.
    As for it's output, it has to be kept at it is, sent bakc to the user."""

    class Config:
        extra = Extra.allow

    def _run(self, question: str):
        # initiate chain and prompts
        # Would be waaay more elegant to have it at init time, but...
        prompt_template = """
        You act as a knowledgeable tutor. Based on some previous [apropos] 
        (a question, a piace of material, or a summary of a tutoring session) 
        and some [relevant documents] generate a tutoring question, 
        that helps in systematically think about the topic at hand and 
        for which the answer is deepening the knowledge of the subject, getting closser to an aswer.
        You should NOT answer the question at hand, just either ask a helping question 
        or confirm if an aswer is correct! This should be [your output].

        [apropos]
        {apropos}

        [relevant_documents]
        {relevant_documents}

        [your output]:
        """

        prompt = PromptTemplate(
            input_variables=["apropos", "relevant_documents"],
            template=prompt_template,
        )

        # do retrieval
        relevant_documents = self.retriever.get_relevant_documents(question)

        # concat the two
        # execute a chain
        llm = ChatOpenAI(model_name=os.environ["CHAT_MODEL"])
        chain = LLMChain(llm=llm, prompt=prompt)

        result = chain.run(apropos=question, relevant_documents=relevant_documents)
        return result
    
    def _arun(self, question: str):
        raise NotImplementedError("This tool does not support async")



class QMLAgent():

    def __init__(self):

        pinecone.init(api_key=os.environ["PINECONE_API_KEY"], environment=os.environ["PINECONE_REGION"])

        index = pinecone.Index(os.environ["INDEX_NAME"])

        embeddings = OpenAIEmbeddings()


        bm25_encoder = BM25Encoder()
        bm25_encoder.load(os.environ["BM25_FILENAME"])

        retriever = PineconeHybridSearchRetriever(
            embeddings=embeddings, 
            sparse_encoder=bm25_encoder, 
            index=index, 
            top_k=os.environ["TOP_K"])

        llm = ChatOpenAI(model_name=os.environ["CHAT_MODEL"], callbacks=[PromptLayerCallbackHandler(pl_tags=["langchain"])])

        math_tool = CalculatorTool()

        random_tool = GetrandomTool()
        random_tool.indexer = index
        random_tool.index_max = index.describe_index_stats()["total_vector_count"] 

        tutoring_tool = TutoringTool()
        tutoring_tool.retriever = retriever

        tools = [
            Tool(
                name="Search",
                func=retriever.get_relevant_documents,
                description=f"You have to use this to search for knowledge about {topic}.",
            ),
            Tool.from_function(
                name="Math calculation",
                func=math_tool._run,
                description=math_tool.description
                #return_direct=False

            ),
            Tool.from_function(
                name="Random document",
                func=random_tool._run,
                description=random_tool.description
                #return_direct=False

            ),
            Tool.from_function(
                name="Tutoring",
                func=tutoring_tool._run,
                description=tutoring_tool.description,
                return_direct=True
            ),
        ]

        memory = ConversationBufferWindowMemory(k=os.environ["MEMORY_LENGTH"], memory_key="chat_history", return_messages=True)

        PREFIX = f"""Assistant is called {assistant_name}, a large language model with a knowledge base about {topic} trained for the {course_name} class at {institution}.

{assistant_name} is designed to be able to assist the students with a range of tasks specifically for the {course_name}, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, {assistant_name} 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.

{assistant_name} is willing to serve the students all the time, but always sticks to the academic context.

It is able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. Additionally, {assistant_name} is able to generate its own text based on the input it receives, allowing it to engage in discussions and provide explanations and descriptions on a wide range of topics.

{assistant_name} is especially helpful in tutoring. If the student explicitly asks for tutoring, {assistant_name} can come up with relevant and interesting questions, pose it to the student and help him/her to discover the answer step by step.

Tutoring is done by 1. recognizing students wish to be asked questions 2. getting a random part of the course material 3. asking some questions from the student 4. hekping him get to the right answer. 

Whether you need help with a specific question or just want to have a conversation about a particular topic, {assistant_name} is here to assist.
"""
# 

        TEMPLATE_TOOL_RESPONSE = """TOOL RESPONSE: 
---------------------
{observation}

USER'S INPUT
--------------------

Okay, so do what I asked for now or answer me if I asked something! If using information obtained from the tools you must mention it explicitly without mentioning the tool names - I have forgotten all TOOL RESPONSES! But if you are talking about code, always explicitly include it, as plain text, no ` marks.  Remember to respond with a markdown code snippet of a json blob with a single action, and NOTHING else, but NEVER use ` marks. There should not be any code blocks, just include code as a normal part of text."""

        self.agent_chain = initialize_agent(
            tools, 
            llm, 
            agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION, 
            verbose=True,
            return_intermediate_steps=False, 
            memory=memory, 
            handle_parsing_errors="Reformat the text to get rid of ` marks, there should be no separation for code blocks.",
            agent_kwargs={"system_message": PREFIX}
        )

        self.agent_chain.agent.template_tool_response = TEMPLATE_TOOL_RESPONSE

    def run(self, question):
        return self.agent_chain.run(question)

if __name__ == '__main__':
    agent = QMLAgent()
    question = "What is eigenvector for the matrix [[1,2,3],[4,5,6],[7,8,9]] raised to the second power?"

    print(agent.run(question))