parser and prompt template

utils.py

@@ -1,17 +1,20 @@
 import os
 import pickle
+import re
+from typing import List, Union
 
-from langchain import LLMChain, OpenAI
-from langchain.agents import ConversationalAgent, AgentExecutor, Tool
-from langchain.memory import ConversationBufferWindowMemory
+import faiss
+from langchain import OpenAI, LLMChain
+from langchain.agents import ConversationalAgent
+from langchain.agents import Tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser
 from langchain.chains import ConversationalRetrievalChain
-from langchain.text_splitter import CharacterTextSplitter
 from langchain.document_loaders import DirectoryLoader, TextLoader, UnstructuredHTMLLoader
-import faiss
-from langchain.vectorstores.faiss import FAISS
 from langchain.embeddings import OpenAIEmbeddings
-
-
+from langchain.memory import ConversationBufferWindowMemory
+from langchain.prompts import BaseChatPromptTemplate
+from langchain.schema import AgentAction, AgentFinish, HumanMessage
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.vectorstores.faiss import FAISS
 
 pickle_file = "open_ai.pkl"
 index_file = "open_ai.index"
@@ -26,6 +29,55 @@ memory = ConversationBufferWindowMemory(memory_key="chat_history")
 
 gpt_3_5_index = None
 
+class CustomOutputParser(AgentOutputParser):
+
+    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
+        # Check if agent replied without using tools
+        if "AI:" in llm_output:
+            return AgentFinish(return_values={"output": llm_output.split("AI:")[-1].strip()},
+                               log=llm_output)
+        # Check if agent should finish
+        if "Final Answer:" in llm_output:
+            return AgentFinish(
+                # Return values is generally always a dictionary with a single `output` key
+                # It is not recommended to try anything else at the moment :)
+                return_values={"output": llm_output.split("Final Answer:")[-1].strip()},
+                log=llm_output,
+            )
+        # Parse out the action and action input
+        regex = r"Action: (.*?)[\n]*Action Input:[\s]*(.*)"
+        match = re.search(regex, llm_output, re.DOTALL)
+        if not match:
+            raise ValueError(f"Could not parse LLM output: `{llm_output}`")
+        action = match.group(1).strip()
+        action_input = match.group(2)
+        # Return the action and action input
+        return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)
+
+# Set up a prompt template
+class CustomPromptTemplate(BaseChatPromptTemplate):
+    # The template to use
+    template: str
+    # The list of tools available
+    tools: List[Tool]
+
+    def format_messages(self, **kwargs) -> str:
+        # Get the intermediate steps (AgentAction, Observation tuples)
+        # Format them in a particular way
+        intermediate_steps = kwargs.pop("intermediate_steps")
+        thoughts = ""
+        for action, observation in intermediate_steps:
+            thoughts += action.log
+            thoughts += f"\nObservation: {observation}\nThought: "
+        # Set the agent_scratchpad variable to that value
+        kwargs["agent_scratchpad"] = thoughts
+        # Create a tools variable from the list of tools provided
+        kwargs["tools"] = "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])
+        # Create a list of tool names for the tools provided
+        kwargs["tool_names"] = ", ".join([tool.name for tool in self.tools])
+        formatted = self.template.format(**kwargs)
+        return [HumanMessage(content=formatted)]
+
 def get_search_index():
     global gpt_3_5_index
     if os.path.isfile(pickle_file) and os.path.isfile(index_file) and os.path.getsize(pickle_file) > 0: