chentianrui
115 lines
8.6 KiB
Python
115 lines
8.6 KiB
Python
import os
|
||
|
||
from llama_index.core import SQLDatabase, SummaryIndex, VectorStoreIndex
|
||
from llama_index.core.indices.struct_store import SQLTableRetrieverQueryEngine
|
||
from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex
|
||
from llama_index.core.settings import Settings
|
||
from llama_index.core.agent import AgentRunner, StructuredPlannerAgent, FunctionCallingAgentWorker, ReActChatFormatter
|
||
from llama_index.core.tools.query_engine import QueryEngineTool
|
||
from sqlalchemy import create_engine, Engine
|
||
from llama_index.core.response_synthesizers.type import ResponseMode
|
||
|
||
from app.engine.loaders.db import makeDescriptionByEngine
|
||
from app.engine.tools import ToolFactory
|
||
from app.engine.index import get_index
|
||
from app.settings import get_node_postprocessors
|
||
|
||
import nest_asyncio
|
||
|
||
nest_asyncio.apply()
|
||
|
||
sql_database = None
|
||
sql_obj_index = None
|
||
|
||
|
||
|
||
def get_chat_engine(filters=None, params=None):
|
||
system_prompt = os.getenv("SYSTEM_PROMPT")
|
||
top_k = int(os.getenv("TOP_K", "3"))
|
||
tools = []
|
||
|
||
global sql_obj_index
|
||
global sql_database
|
||
if sql_obj_index is None:
|
||
sqlengine = create_engine(os.getenv("SQL_DATABASE_URL", ""))
|
||
sql_database = SQLDatabase(sqlengine)
|
||
table_schema_objs = makeDescriptionByEngine(sql_database)
|
||
table_node_mapping = SQLTableNodeMapping(sql_database)
|
||
|
||
sql_obj_index = ObjectIndex.from_objects(
|
||
table_schema_objs,
|
||
table_node_mapping,
|
||
index_cls=VectorStoreIndex,
|
||
)
|
||
|
||
# 创建SQL查询工具
|
||
sql_query_engine = SQLTableRetrieverQueryEngine(sql_database,
|
||
sql_obj_index.as_retriever(similarity_top_k=top_k),
|
||
verbose=True,)
|
||
sql_query_tool = QueryEngineTool.from_defaults(query_engine=sql_query_engine,
|
||
name="zjdata_query_tool",
|
||
description="来源于一个由博微公司电力造价软件编制的造价工程文件。该文件以多张表格的形式存储存储了整个工程的全部数据内容。适用于以详细的自然语言查询表格数据方式查询造价工程各项具体属性、费用的数值。请先使用“zj_query_tool”无法解决才使用本工具")
|
||
|
||
# Add query tool if index exists
|
||
index = get_index()
|
||
if index is not None:
|
||
summary_index = SummaryIndex(index.vector_store.get_nodes(node_ids=None))
|
||
summary_query_engine = summary_index.as_query_engine()
|
||
summary_query_tool = QueryEngineTool.from_defaults( query_engine=summary_query_engine, name="summary_query_tool",
|
||
description="适用于任何需要进行全面总结、概括的要求。",
|
||
#description="适用于任何需要对所有内容进行全面总结的请求。有关电力造价领域更具体部分的问题，请使用zj_query_engine_tool",
|
||
)
|
||
|
||
# 创建向量检索查询工具
|
||
postprocess = get_node_postprocessors()
|
||
query_engine = index.as_query_engine(
|
||
similarity_top_k=top_k, filters=filters,
|
||
node_postprocessors=postprocess,
|
||
use_async=True,
|
||
streaming=True,
|
||
# response_mode=ResponseMode.TREE_SUMMARIZE,
|
||
)
|
||
query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine, name="zj_query_tool",
|
||
description="由博微公司编制的关于电力造价知识、电力造价编制软件知识和造价工程文件结构的知识库。适用于查询电力领域、电力造价领域、博微、博微电力、博微造价等业务等内容。如果本知识库没有直接答案但有解决思路的可以返回解决办法后建议使用“zjdata_query_tool”工具。如果你不知道答案，就说你不知道，不要编造答案。",
|
||
)
|
||
|
||
tools.append(summary_query_tool)
|
||
tools.append(query_engine_tool)
|
||
#tools.append(sql_query_tool)
|
||
|
||
# Add additional tools
|
||
tools += ToolFactory.from_env()
|
||
|
||
# agentrunner = StructuredPlannerAgent.from_llm(
|
||
# llm=Settings.llm,
|
||
# tools=tools,
|
||
# system_prompt=system_prompt,
|
||
# verbose=True,
|
||
# )
|
||
|
||
# prompts = agentrunner.agent_worker._get_prompts()
|
||
# prompts["system_prompt"].template = """您的设计旨在帮助完成各种任务，从回答问题到提供其他类型分析的摘要。\n\n##工具\n\n你可以访问各种工具。你有责任按照你认为合适的顺序使用这些工具来完成当前的任务。\n这可能需要将任务分解为子任务，并使用不同的工具来完成每个子任务。\n\n你可以访问以下工具：\n{tool_desc}\n\n\n##输出格式\n\n请用与问题相同的语言回答，并使用以下格式：\n\n```\nThought: 用户当前的语言是：(user's language)。我需要使用工具来帮助我回答问题。\nAction: 如果使用工具,则为工具名称(one of {tool_names})。\nAction Input: 输入给工具的内容，使用JSON格式表示kwargs（例如{{\"input\": \"hello world\", \"num_beams\": 5}}）\n```\n\n请始终以Thought开始。\n\n切勿用Markdown代码标记包围你的响应。如果需要，可以在响应中使用代码标记。\n\n请为Action Input使用有效的JSON格式。不要这样做{{\'input\': \'hello world\', \'num_beams\': 5}}。\n\n如果使用此格式，用户将以下面的格式进行回应：\n\n```\nObservation: 工具响应\n```\n\n你应该继续重复上述格式，直到你有足够的信息来回答问题而无需使用更多工具。此时，你必须使用以下两种格式之一进行回答：\n\n```\nThought: 我可以不用任何工具来回答。我将使用用户的语言来回答。\nAnswer: [你的答案（与用户问题相同的语言）]\n```\n\n```\nThought: 我无法使用提供的工具回答问题。\nAnswer: [你的答案（与用户问题相同的语言）]\n```\n\n## 当前对话\n\n以下是当前对话，由人类和助手的消息交替组成。\n"""
|
||
# agentrunner.agent_worker.update_prompts(prompts)
|
||
|
||
prefix_messages = ("""您的设计旨在帮助完成各种任务，从回答问题到提供其他类型分析的摘要。\n\n##工具\n\n你可以访问各种工具。你有责任按照你认为合适的顺序使用这些工具来完成当前的任务。\n这可能需要将任务分解为子任务，并使用不同的工具来完成每个子任务。\n\n你可以访问以下工具：\n{tool_desc}\n\n\n##输出格式\n\n请用与问题相同的语言回答，并使用以下格式：\n\n```\nThought: 用户当前的语言是：(user's language)。我需要使用工具来帮助我回答问题。\nAction: 如果使用工具,则为工具名称(one of {tool_names})。\nAction Input: 输入给工具的内容，使用JSON格式表示kwargs（例如{{\"input\": \"hello world\", \"num_beams\": 5}}）\n```\n\n请始终以Thought开始。\n\n切勿用Markdown代码标记包围你的响应。如果需要，可以在响应中使用代码标记。\n\n请为Action Input使用有效的JSON格式。不要这样做{{\'input\': \'hello world\', \'num_beams\': 5}}。\n\n如果使用此格式，用户将以下面的格式进行回应：\n\n```\nObservation: 工具响应\n```\n\n你应该继续重复上述格式，直到你有足够的信息来回答问题而无需使用更多工具。此时，你必须使用以下两种格式之一进行回答：\n\n```\nThought: 我可以不用任何工具来回答。我将使用用户的语言来回答。\nAnswer: [你的答案（与用户问题相同的语言）]\n```\n\n```\nThought: 我无法使用提供的工具回答问题。\nAnswer: [你的答案（与用户问题相同的语言）]\n```\n\n##如果从工具中得到的回应是Empty Response，那么只需要回答“我不知道”，不需要额外回答别的内容。## 当前对话\n\n以下是当前对话，由人类和助手的消息交替组成。\n""")
|
||
|
||
react_chat_formatter = ReActChatFormatter.from_defaults(prefix_messages)
|
||
|
||
agentrunner = AgentRunner.from_llm(
|
||
llm=Settings.llm,
|
||
tools=tools,
|
||
react_chat_formatter=react_chat_formatter,
|
||
system_prompt=system_prompt,
|
||
verbose=True,
|
||
)
|
||
return agentrunner
|
||
|
||
|
||
# create the function calling worker for reasoning
|
||
# worker = FunctionCallingAgentWorker.from_tools(
|
||
# tools, verbose=True
|
||
# )
|
||
#
|
||
# # wrap the worker in the top-level planner
|
||
# return StructuredPlannerAgent(worker, tools)
|