82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
from langchain.agents import load_tools
|
|
from langchain.agents import ZeroShotAgent, Tool, AgentExecutor
|
|
from langchain.memory import ConversationBufferMemory
|
|
from langchain.agents import initialize_agent
|
|
from langchain.agents import AgentType
|
|
from langchain import LLMChain
|
|
|
|
|
|
class CompletionMeta(type):
|
|
"""
|
|
The meta class for completion interface
|
|
"""
|
|
|
|
tools = None
|
|
|
|
class Completion():
|
|
def __init__(self):
|
|
self.prefix = None
|
|
self.suffix = None
|
|
self.promptVars = None
|
|
|
|
def complete(self, message: str) -> str:
|
|
"""
|
|
Perform a text completion using the language model
|
|
"""
|
|
pass
|
|
|
|
def getModelName(self) -> str:
|
|
"""
|
|
Return the model name
|
|
"""
|
|
pass
|
|
|
|
def getChain(self, llm):
|
|
pass
|
|
|
|
def getAgent(self, llm):
|
|
# Load tools
|
|
tools = load_tools(["serpapi", "llm-math"], llm=llm)
|
|
|
|
# Build a prompt
|
|
prompt = ZeroShotAgent.create_prompt(
|
|
tools=tools,
|
|
prefix=self.getPromptPrefix(),
|
|
suffix=self.getPromptSuffix(),
|
|
input_variables=self.getPromptVariables()
|
|
)
|
|
memory = ConversationBufferMemory(memory_key="chat_history")
|
|
|
|
# Build the LLM Chain
|
|
llm_chain = LLMChain(llm=llm, prompt=prompt)
|
|
agent = ZeroShotAgent(llm_chain=llm_chain,
|
|
tools=tools,
|
|
verbose=True)
|
|
self.agent_chain = AgentExecutor.from_agent_and_tools(agent=agent,
|
|
tools=tools,
|
|
verbose=True,
|
|
memory=memory)
|
|
|
|
return self.agent_chain
|
|
|
|
def getPromptVariables(self):
|
|
if self.promptVars is None:
|
|
self.promptVars = ["input", "chat_history", "agent_scratchpad"]
|
|
return self.promptVars
|
|
|
|
def getPromptPrefix(self):
|
|
if self.prefix is None:
|
|
self.prefix = """Have a conversation with a human, answering the following
|
|
questions as best you can. You have access to the following tools:"""
|
|
return self.prefix
|
|
|
|
def getPromptSuffix(self):
|
|
if self.suffix is None:
|
|
self.suffix = """Begin!
|
|
{chat_history}
|
|
Question: {input}
|
|
{agent_scratchpad}"""
|
|
return self.suffix
|
|
|
|
|