Class AgentHelper

gpt/agent.py:14–102

Add agent to help docGPT can be perfonm better.

class AgentHelper:
    """Add agent to help docGPT can be perfonm better."""
    def __init__(self) -> None:
        self._llm = None
        self.agent_ = None
        self.tools = []

    @property
    def llm(self):
        return self._llm

    @llm.setter
    def llm(self, llm) -> None:
        self._llm = llm

    @property
    def get_calculate_chain(self) -> Tool:
        from langchain import LLMMathChain

        llm_math_chain = LLMMathChain.from_llm(llm=self.llm, verbose=True)
        tool = Tool(
            name='Calculator',
            func=llm_math_chain.run,
            description='useful for when you need to answer questions about math'
        )
        return tool

    @property
    def get_searp_chain(self) -> Tool:
        from langchain import SerpAPIWrapper

        search = SerpAPIWrapper()
        tool = Tool(
            name='Search',
            func=search.run,
            description='useful for when you need to answer questions about current events'
        )
        return tool

    def create_doc_chat(self, docGPT) -> Tool:
        """Add a custom docGPT tool"""
        tool = Tool(
            name='DocumentGPT',
            func=docGPT.run,
            description="""
            useful for when you need to answer questions from the context of PDF
            """
        )
        return tool

    def create_llm_chain(self) -> Tool:
        """Add a llm tool"""
        prompt = PromptTemplate(
            input_variables = ['query'],
            template = '{query}'
        )
        llm_chain = LLMChain(llm=self.llm, prompt=prompt)