Add: LLM model in agent

Author: Lin-jun-xiang

agent/agent.py

@@ -5,8 +5,9 @@
 from langchain import LLMMathChain, SerpAPIWrapper
 from langchain.agents import AgentType, Tool, initialize_agent
 from langchain.callbacks import get_openai_callback
+from langchain.chains import LLMChain
 from langchain.llms import OpenAI
-
+from langchain.prompts import PromptTemplate
 
 openai.api_key = os.getenv('OPENAI_API_KEY')
 os.environ['SERPAPI_API_KEY'] = os.getenv('SERPAPI_API_KEY')
@@ -51,6 +52,20 @@ def create_doc_chat(self, docGPT) -> Tool:
         )
         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)
+
+        tool = Tool(
+            name='LLM',
+            func=llm_chain.run,
+            description='useful for general purpose queries and logic'
+        )
+        return tool
     def initialize(self, tools):
         for tool in tools:
             if isinstance(tool, Tool):
@@ -66,6 +81,9 @@ def initialize(self, tools):
     def query(self, query: str) -> Optional[str]:
         response = None
         with get_openai_callback() as callback:
+            # TODO: The true result will hide in 'Observation'
+            # https://github.com/hwchase17/langchain/issues/4916
+            # https://python.langchain.com/docs/modules/agents/how_to/intermediate_steps
             response = self.agent_.run(query)
             print(callback)
         return response

app.py

@@ -96,7 +96,10 @@ def load_api_key() -> None:
             docGPT.create_qa_chain(
                 chain_type='refine',
             )
+
             docGPT_tool = agent_.create_doc_chat(docGPT)
+            calculate_tool = agent_.get_calculate_chain
+            llm_tool = agent_.create_llm_chain()
 
         except Exception as e:
             print(e)
@@ -107,11 +110,10 @@ def load_api_key() -> None:
             print(e)
 
         try:
-            calculate_tool = agent_.get_calculate_chain
-
             tools = [
                 docGPT_tool,
-                search_tool
+                search_tool,
+                llm_tool
             ]
             agent_.initialize(tools)
         except Exception as e: