Merge pull request #27 from FacerAin/feat/agent

feat: update agent model to OpenAIFunctionsAgent

app/agent/agent.py

@@ -1,13 +1,14 @@
 import datetime
 
-from langchain.agents import AgentExecutor, AgentType, LLMSingleActionAgent, Tool, initialize_agent
+from langchain.agents import AgentExecutor, LLMSingleActionAgent, OpenAIFunctionsAgent, Tool
 from langchain.cache import InMemoryCache
 from langchain.chains import LLMChain
 from langchain.chat_models import ChatOpenAI
 from langchain.globals import set_llm_cache
+from langchain_core.messages import SystemMessage
 
 from app.agent.parser import CustomAgentOutputParser
-from app.agent.prompts import AgentPromptTemplate, agent_prompt_template, retriever_prompt_template
+from app.agent.prompts import retriever_prompt_template, system_message
 from app.agent.retriever import PineconeRetriever
 from app.core.config import settings
 
@@ -32,29 +33,23 @@ def __init__(self):
                 func=simple_meal_info,
                 description="If a user is looking for campus cafeterial menu information, use this information.",
             ),
+            Tool(
+                name="today_date",
+                func=get_today_date,
+                description="If you’re curious about today’s date, you can use this. It returns today’s date as a string in the format of YYYY-mm-dd.",
+            ),
         ]
-        self.agent_prompt = AgentPromptTemplate(
-            template=agent_prompt_template,
-            tools=self.tools,
-            input_variables=["input", "intermediate_steps"],
-        )
-        self.agent = initialize_agent(
-            tools=self.tools, llm=self.llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=True, max_iterations=3
-        )
+
+        self.agent_prompt = OpenAIFunctionsAgent.create_prompt(system_message=SystemMessage(content=system_message))
         self.output_parser = CustomAgentOutputParser()
-        llm_chain = LLMChain(llm=self.llm, prompt=self.agent_prompt)
-        tool_names = [tool.name for tool in self.tools]
-        self.agent = LLMSingleActionAgent(
-            llm_chain=llm_chain, output_parser=self.output_parser, stop=["\nObservation:"], allowed_tools=tool_names
-        )
+        self.agent = OpenAIFunctionsAgent(llm=self.llm, tools=self.tools, prompt=self.agent_prompt)
 
         self.executor = AgentExecutor.from_agent_and_tools(
             agent=self.agent, tools=self.tools, verbose=True, max_iterations=2
         )
 
     def run(self, query):
         response = self.executor.run(query)
-        print(response)
         return response
 
 
@@ -66,6 +61,11 @@ def simple_meal_info(query):
 """
 
 
+def get_today_date(query):
+    now = datetime.datetime.now()
+    return str(now.strftime("%Y-%m-%d"))
+
+
 class RetrieverAgent:
     def __init__(self, index_name: str = "khugpt") -> None:
         self.llm = ChatOpenAI(model_name="gpt-4-1106-preview", temperature=0, openai_api_key=settings.OPENAI_API_KEY)

app/agent/prompts.py

@@ -3,6 +3,11 @@
 from langchain.agents import Tool
 from langchain.prompts import PromptTemplate, StringPromptTemplate
 
+system_message = """You are a helpful assistant for Kyung Hee University students.
+Answer the following questions as best you can. If a page_url is provided in the document, please also provide a link to the related page.
+Remember to speak in a korean when giving your final answer.
+"""
+
 agent_prompt_template = """You are a helpful assistant for Kyung Hee University students.
 Answer the following questions as best you can. If a page_url is provided in the document, please also provide a link to the related page.
 You have access to the following tools: