254 lines
14 KiB
Python
254 lines
14 KiB
Python
|
from datetime import datetime
|
|||
|
|
import uuid
|
|||
|
|
from langchain.memory import ConversationBufferMemory
|
|||
|
|
from server.agent.agent import Agent
|
|||
|
|
from server.agent.custom_agent.ChatGLM3Agent import initialize_glm3_agent
|
|||
|
|
from server.agent.tools_select import tools, tool_names, search_tool_names
|
|||
|
|
from server.agent.callbacks import CustomAsyncIteratorCallbackHandler, Status
|
|||
|
|
from langchain.agents import LLMSingleActionAgent, AgentExecutor
|
|||
|
|
from server.agent.custom_template import CustomOutputParser, CustomPromptTemplate
|
|||
|
|
from fastapi import Body
|
|||
|
|
from sse_starlette.sse import EventSourceResponse
|
|||
|
|
from configs import LLM_MODELS, TEMPERATURE, HISTORY_LEN, Agent_MODEL
|
|||
|
|
from server.chat import utils
|
|||
|
|
from server.chat.knowledge_base_name import KnowledgeBase
|
|||
|
|
from server.utils import replace_variables, wrap_done, get_ChatOpenAI, get_prompt_template
|
|||
|
|
from langchain.chains import LLMChain
|
|||
|
|
from typing import AsyncIterable, Optional
|
|||
|
|
from server.agent.tools import rag_search
|
|||
|
|
import asyncio
|
|||
|
|
from typing import List
|
|||
|
|
from server.chat.utils import History, remove_after_and_including, remove_after_and_includings
|
|||
|
|
import json
|
|||
|
|
from server.agent import model_container
|
|||
|
|
from server.knowledge_base.kb_service.base import get_kb_details
|
|||
|
|
import ast
|
|||
|
|
import re
|
|||
|
|
import asyncio
|
|||
|
|
from concurrent.futures import ThreadPoolExecutor
|
|||
|
|
from server.chat.policy_fun_iast import get_llm_model_response, get_llm_model_response_async, get_llm_model_response_stream_openai
|
|||
|
|
from configs import kb_config
|
|||
|
|
from configs.basic_config import *
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
_executor: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=8)
|
|||
|
|
|
|||
|
|
async def run_sync(func, /, *args, **kwargs):
|
|||
|
|
loop = asyncio.get_running_loop()
|
|||
|
|
return await loop.run_in_executor(_executor, lambda: func(*args, **kwargs))
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def agent_write_test(
|
|||
|
|
user_prompt_name: Optional[str] = Body(None, description="用户输入", examples=[""]),
|
|||
|
|
style: Optional[str] = Body("默认风格", description="语言风格", examples=[""]),
|
|||
|
|
query: str = Body(..., description="用户输入", examples=["恼羞成怒"]),
|
|||
|
|
think_content: str = Body(..., description="思考过程", examples=[""]),
|
|||
|
|
uuid: str=Body(..., description="用户输入", examples=["恼羞成怒"]),
|
|||
|
|
history: List[History] = Body([],
|
|||
|
|
description="历史对话",
|
|||
|
|
examples=[[
|
|||
|
|
{"role": "user",
|
|||
|
|
"content": "请使用天气查询工具查询今天北京天气"},
|
|||
|
|
{"role": "assistant",
|
|||
|
|
"content": "今天是2024年3月22日，受冷空气影响，白天有3、4级偏北风，阵风6、7"
|
|||
|
|
"级，西部山区阵风相对明显，局地伴有扬沙。白天晴，局地有扬沙,偏北风,1级转3、4级，阵风6、7级,"
|
|||
|
|
"最高气温22℃。夜间晴间多云,偏北风,1、2级,最低气温6℃。"}]]
|
|||
|
|
),
|
|||
|
|
stream: bool = Body(False, description="流式输出"),
|
|||
|
|
model_name: str = Body(LLM_MODELS[0], description="LLM 模型名称。"),
|
|||
|
|
temperature: float = Body(0.1, description="LLM 采样温度", ge=0.0, le=1.0),
|
|||
|
|
max_tokens: Optional[int] = Body(None, description="限制LLM生成Token数量，默认None代表模型最大值"),
|
|||
|
|
prompt_name: str = Body("default",
|
|||
|
|
description="使用的prompt模板名称(在configs/prompt_config.py中配置)"),
|
|||
|
|
) -> AsyncIterable[str]:
|
|||
|
|
agent = Agent()
|
|||
|
|
agent.step = "初始步骤"
|
|||
|
|
finish_tools = "已经调用过的工具名称："
|
|||
|
|
if user_prompt_name == "complete_outline":
|
|||
|
|
user_prompt = get_prompt_template("llm_chat",user_prompt_name+"_with_history")
|
|||
|
|
user_prompt = replace_variables(user_prompt, replace_content=style, replace_param="{style}")
|
|||
|
|
else:
|
|||
|
|
user_prompt = get_prompt_template("llm_chat",user_prompt_name)
|
|||
|
|
user_prompt = user_prompt.replace("{{time}}",datetime.now().strftime("%Y%m%d")).replace("{{input}}",query)
|
|||
|
|
# 假设 Tool 有 .name 和 .func 属性
|
|||
|
|
tool_map = {tool.name: tool.func for tool in tools}
|
|||
|
|
if history:
|
|||
|
|
history_detail = str(history)
|
|||
|
|
else:
|
|||
|
|
history_detail = ""
|
|||
|
|
all_tool = "\n".join([f'【工具名称】{tool.name}: 【工具描述】{tool.description}' for tool in tools])
|
|||
|
|
rewrite = 0
|
|||
|
|
wrong_num = 0
|
|||
|
|
while not agent.step == "answer":
|
|||
|
|
if rewrite == 0:
|
|||
|
|
step_res = await run_sync(
|
|||
|
|
get_llm_model_response,
|
|||
|
|
strategy_name="query rewrite",
|
|||
|
|
llm_model_name=LLM_MODELS[0],
|
|||
|
|
template_prompt_name="get_next_write_tip",
|
|||
|
|
prompt_param_dict={
|
|||
|
|
"time": datetime.now().strftime("%Y%m%d"),
|
|||
|
|
"step": agent.step,
|
|||
|
|
"question": query,
|
|||
|
|
"user_prompt": user_prompt,
|
|||
|
|
"history": history_detail,
|
|||
|
|
"tools": all_tool,
|
|||
|
|
"res": agent.res,
|
|||
|
|
"finish_tools": finish_tools,
|
|||
|
|
"think_content": think_content,
|
|||
|
|
},
|
|||
|
|
temperature=0.01,
|
|||
|
|
max_tokens=512,
|
|||
|
|
)
|
|||
|
|
next_step_pattern = r'<step>(.*?)</step>'
|
|||
|
|
next_step_content = re.search(next_step_pattern, step_res, re.DOTALL)
|
|||
|
|
agent.step = next_step_content.group(1).strip() if next_step_content else ""
|
|||
|
|
agent.res += (
|
|||
|
|
step_res.replace("<step>", "")
|
|||
|
|
.replace("</step>", "")
|
|||
|
|
.replace(agent.step, "")
|
|||
|
|
)
|
|||
|
|
match agent.step:
|
|||
|
|
case "thinking":
|
|||
|
|
thinking = await run_sync(
|
|||
|
|
get_llm_model_response,
|
|||
|
|
strategy_name="query rewrite",
|
|||
|
|
llm_model_name=LLM_MODELS[0],
|
|||
|
|
template_prompt_name="agent_write_think",
|
|||
|
|
prompt_param_dict={
|
|||
|
|
"time": datetime.now().strftime("%Y%m%d"),
|
|||
|
|
"user_prompt": user_prompt,
|
|||
|
|
"think_content": think_content,
|
|||
|
|
"input": query,
|
|||
|
|
"history": history_detail,
|
|||
|
|
"res": agent.res,
|
|||
|
|
"tools": all_tool,
|
|||
|
|
"finish_tools": finish_tools,
|
|||
|
|
},
|
|||
|
|
temperature=0.01,
|
|||
|
|
max_tokens=512,
|
|||
|
|
)
|
|||
|
|
agent.res = agent.res+"<step>thinking</step>"+thinking
|
|||
|
|
print("thinking")
|
|||
|
|
case "select_tool":
|
|||
|
|
tool_desc = await run_sync(
|
|||
|
|
get_llm_model_response,
|
|||
|
|
strategy_name="query rewrite",
|
|||
|
|
llm_model_name=LLM_MODELS[0],
|
|||
|
|
template_prompt_name="tool_write_select",
|
|||
|
|
prompt_param_dict={
|
|||
|
|
"time": datetime.now().strftime("%Y%m%d"),
|
|||
|
|
"user_prompt": user_prompt,
|
|||
|
|
"think_content": think_content,
|
|||
|
|
"input": query,
|
|||
|
|
"history": history_detail,
|
|||
|
|
"res": agent.res,
|
|||
|
|
"tools": tools,
|
|||
|
|
"finish_tools": finish_tools
|
|||
|
|
},
|
|||
|
|
temperature=0.01,
|
|||
|
|
max_tokens=512,
|
|||
|
|
)
|
|||
|
|
agent.res += "<step>select_tool</step>"
|
|||
|
|
try:
|
|||
|
|
try:
|
|||
|
|
tool_name_pattern = r'<tool>(.*?)</tool>'
|
|||
|
|
toolname = re.search(tool_name_pattern, tool_desc.replace("\n","")).group(1).strip()
|
|||
|
|
|
|||
|
|
tool_input_pattern = r'<tool_input>(.*?)</tool_input>'
|
|||
|
|
toolinput=re.search(tool_input_pattern, tool_desc.replace("\n","")).group(1).strip()
|
|||
|
|
except Exception as e:
|
|||
|
|
if wrong_num > 1:
|
|||
|
|
tool_desc += f"请重新修正。以修正失败{wrong_num}次"
|
|||
|
|
tool_desc = await run_sync(
|
|||
|
|
get_llm_model_response,
|
|||
|
|
strategy_name="query rewrite",
|
|||
|
|
llm_model_name=LLM_MODELS[0],
|
|||
|
|
template_prompt_name="agent_rewrite",
|
|||
|
|
prompt_param_dict={"input": tool_desc, "format": all_tool},
|
|||
|
|
temperature=0.01,
|
|||
|
|
max_tokens=512,
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
tool_name_pattern = r'<tool>(.*?)</tool>'
|
|||
|
|
toolname = re.search(tool_name_pattern, tool_desc.replace("\n","")).group(1).strip()
|
|||
|
|
|
|||
|
|
tool_input_pattern = r'<tool_input>(.*?)</tool_input>'
|
|||
|
|
toolinput=re.search(tool_input_pattern, tool_desc.replace("\n","")).group(1).strip()
|
|||
|
|
except Exception as e:
|
|||
|
|
print(e)
|
|||
|
|
wrong_num += 1
|
|||
|
|
rewrite = 1
|
|||
|
|
if wrong_num > 3:
|
|||
|
|
rewrite = 0
|
|||
|
|
finish_tools += tool_desc
|
|||
|
|
agent.res+=f"使用{tool_desc}工具出现了异常，请使用其他工具或用自身能力回答，禁止虚拟链接"
|
|||
|
|
finish_tools = finish_tools+ toolname +","
|
|||
|
|
rewrite = 0
|
|||
|
|
agent.res+="使用"+toolname+"工具的输入：【"+toolinput+"】"
|
|||
|
|
toolinput+="{\"uuid\":\""+uuid+"\"}"
|
|||
|
|
tool_func = tool_map[toolname]
|
|||
|
|
# ---------- 调用工具（阻塞 → 线程池） ----------
|
|||
|
|
if asyncio.iscoroutinefunction(tool_func):
|
|||
|
|
result = await tool_func(toolinput)
|
|||
|
|
else:
|
|||
|
|
result = await run_sync(tool_func, toolinput)
|
|||
|
|
if result == None and toolname == "统计数据查询":
|
|||
|
|
agent.res+="使用"+toolname+"工具没有搜索到结果,推荐用联网思索替换统计数据查询"
|
|||
|
|
else:
|
|||
|
|
agent.res+="使用"+toolname+"工具的结果：【"+result+"】"
|
|||
|
|
rewrite = 0
|
|||
|
|
if toolname == "知识库联想":
|
|||
|
|
source = utils.get_shared_variable(uuid)
|
|||
|
|
source_docs = source["source_docs"]
|
|||
|
|
try:
|
|||
|
|
docs_string = "\n"+"\n".join(f"{str(doc)}\n" for doc in source_docs)
|
|||
|
|
except:
|
|||
|
|
print("知识库联想存在异常")
|
|||
|
|
pass
|
|||
|
|
yield json.dumps({"docs": docs_string}, ensure_ascii=False)
|
|||
|
|
if toolname == "联网思索":
|
|||
|
|
source = utils.get_shared_variable(uuid)
|
|||
|
|
source_docs = source["source_docs"]
|
|||
|
|
try:
|
|||
|
|
docs_string = "\n"+"\n".join(f"{str(doc)}\n" for doc in source_docs)
|
|||
|
|
except:
|
|||
|
|
print("联网思索存在异常")
|
|||
|
|
pass
|
|||
|
|
yield json.dumps({"docs": docs_string}, ensure_ascii=False)
|
|||
|
|
# if toolname == "天气工具":
|
|||
|
|
# source = utils.get_shared_variable(uuid)
|
|||
|
|
# source_docs = source["source_docs"]
|
|||
|
|
# try:
|
|||
|
|
# docs_string = "\n"+"\n".join(f"{str(doc)}\n" for doc in source_docs)
|
|||
|
|
# except:
|
|||
|
|
# print("知识库联想存在异常")
|
|||
|
|
# pass
|
|||
|
|
# yield json.dumps({"docs": docs_string}, ensure_ascii=False)
|
|||
|
|
print("结果：", result)
|
|||
|
|
print("select_tool")
|
|||
|
|
rewrite = 0
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
rewrite = 1
|
|||
|
|
wrong_num +=0.5
|
|||
|
|
if wrong_num > 3:
|
|||
|
|
rewrite = 0
|
|||
|
|
finish_tools += tool_desc
|
|||
|
|
agent.res+=f"使用{tool_desc}工具出现了异常，请使用其他工具或用自身能力回答，禁止虚拟链接"
|
|||
|
|
await asyncio.sleep(0)
|
|||
|
|
|
|||
|
|
async for chunk in get_llm_model_response_stream_openai(
|
|||
|
|
type=1,
|
|||
|
|
strategy_name="query rewrite",
|
|||
|
|
llm_model_name=LLM_MODELS[0],
|
|||
|
|
template_prompt_name="agent_write_answer",
|
|||
|
|
prompt_param_dict={"time": datetime.now().strftime("%Y%m%d"),"user_prompt":user_prompt,"think_content":think_content,"input":query,"history":history_detail,"res":agent.res,"tools":all_tool,"finish_tools":finish_tools},
|
|||
|
|
temperature=0.01,
|
|||
|
|
max_tokens=20000
|
|||
|
|
):
|
|||
|
|
yield json.dumps(
|
|||
|
|
{"answer": chunk},
|
|||
|
|
ensure_ascii=False)
|