File size: 5,060 Bytes
f7c8212 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
from langchain.agents import Tool
from langchain.tools import BaseTool
from langchain import PromptTemplate, LLMChain
from agent.custom_search import DeepSearch
from langchain.agents import BaseSingleActionAgent, AgentOutputParser, LLMSingleActionAgent, AgentExecutor
from typing import List, Tuple, Any, Union, Optional, Type
from langchain.schema import AgentAction, AgentFinish
from langchain.prompts import StringPromptTemplate
from langchain.callbacks.manager import CallbackManagerForToolRun
from langchain.base_language import BaseLanguageModel
import re
agent_template = """
你现在是一个{role}。这里是一些已知信息：
{related_content}
{background_infomation}
{question_guide}：{input}
{answer_format}
"""
class CustomPromptTemplate(StringPromptTemplate):
template: str
tools: List[Tool]
def format(self, **kwargs) -> str:
intermediate_steps = kwargs.pop("intermediate_steps")
# 没有互联网查询信息
if len(intermediate_steps) == 0:
background_infomation = "\n"
role = "傻瓜机器人"
question_guide = "我现在有一个问题"
answer_format = "如果你知道答案，请直接给出你的回答！如果你不知道答案，请你只回答\"DeepSearch('搜索词')\"，并将'搜索词'替换为你认为需要搜索的关键词，除此之外不要回答其他任何内容。\n\n下面请回答我上面提出的问题！"
# 返回了背景信息
else:
# 根据 intermediate_steps 中的 AgentAction 拼装 background_infomation
background_infomation = "\n\n你还有这些已知信息作为参考：\n\n"
action, observation = intermediate_steps[0]
background_infomation += f"{observation}\n"
role = "聪明的 AI 助手"
question_guide = "请根据这些已知信息回答我的问题"
answer_format = ""
kwargs["background_infomation"] = background_infomation
kwargs["role"] = role
kwargs["question_guide"] = question_guide
kwargs["answer_format"] = answer_format
return self.template.format(**kwargs)
class CustomSearchTool(BaseTool):
name: str = "DeepSearch"
description: str = ""
def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None):
return DeepSearch.search(query = query)
async def _arun(self, query: str):
raise NotImplementedError("DeepSearch does not support async")
class CustomAgent(BaseSingleActionAgent):
@property
def input_keys(self):
return ["input"]
def plan(self, intermedate_steps: List[Tuple[AgentAction, str]],
**kwargs: Any) -> Union[AgentAction, AgentFinish]:
return AgentAction(tool="DeepSearch", tool_input=kwargs["input"], log="")
class CustomOutputParser(AgentOutputParser):
def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
# group1 = 调用函数名字
# group2 = 传入参数
match = re.match(r'^[\s\w]*(DeepSearch)\(([^\)]+)\)', llm_output, re.DOTALL)
print(match)
# 如果 llm 没有返回 DeepSearch() 则认为直接结束指令
if not match:
return AgentFinish(
return_values={"output": llm_output.strip()},
log=llm_output,
)
# 否则的话都认为需要调用 Tool
else:
action = match.group(1).strip()
action_input = match.group(2).strip()
return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)
class DeepAgent:
tool_name: str = "DeepSearch"
agent_executor: any
tools: List[Tool]
llm_chain: any
def query(self, related_content: str = "", query: str = ""):
tool_name = self.tool_name
result = self.agent_executor.run(related_content=related_content, input=query ,tool_name=self.tool_name)
return result
def __init__(self, llm: BaseLanguageModel, **kwargs):
tools = [
Tool.from_function(
func=DeepSearch.search,
name="DeepSearch",
description=""
)
]
self.tools = tools
tool_names = [tool.name for tool in tools]
output_parser = CustomOutputParser()
prompt = CustomPromptTemplate(template=agent_template,
tools=tools,
input_variables=["related_content","tool_name", "input", "intermediate_steps"])
llm_chain = LLMChain(llm=llm, prompt=prompt)
self.llm_chain = llm_chain
agent = LLMSingleActionAgent(
llm_chain=llm_chain,
output_parser=output_parser,
stop=["\nObservation:"],
allowed_tools=tool_names
)
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)
self.agent_executor = agent_executor
|