#!/usr/bin/env python # -*- encoding:utf-8 -*- """ 微信企业后台 接口. """ import xml.etree.cElementTree as ET from bizwechat import WXBizMsgCrypt from http import HTTPStatus import fastapi import logging import config import asyncio import aiohttp from sqlcode.sql_agent import create_sql_agent, remove_markdown_code_block from langchain_community.agent_toolkits import SQLDatabaseToolkit from langchain_community.utilities import SQLDatabase from langchain.agents.agent_types import AgentType from langchain_core.tools import StructuredTool from langchain.prompts import PromptTemplate from langchain_chroma import Chroma from langchain_community.embeddings import DashScopeEmbeddings from sqlcode.utils import parse, format_docs from langchain_core.runnables import RunnablePassthrough from sqlcode.multi_agent import create_sql_graph app = fastapi.FastAPI() logger = logging.getLogger('sqlcode') def get_wxcpt(): ''' 创建微信企业后台接口的加密解密对象 ''' wxbiz_config = config.bizwechat_config() return WXBizMsgCrypt(sToken=wxbiz_config.token, sEncodingAESKey=wxbiz_config.aes_key, sReceiveId=wxbiz_config.corp_id) @app.get('/', include_in_schema=False) async def verify_url(request:fastapi.Request): ''' 验证微信公众号的请求URL的合法性 ''' signature = request.query_params.get('msg_signature') timestamp = request.query_params.get('timestamp') nonce = request.query_params.get('nonce') echostr = request.query_params.get('echostr') if not signature or not timestamp or not nonce or not echostr: logger.error('verify_url failed, missing parameters') return fastapi.Response('', HTTPStatus.BAD_REQUEST) code, echostr = get_wxcpt().VerifyURL(signature, timestamp, nonce, echostr) if code == 0: logger.info('verify_url success, echostr: %s', echostr) return fastapi.Response(echostr, HTTPStatus.OK) else: logger.error('verify_url failed, error code: %s', code) return fastapi.Response('', HTTPStatus.BAD_REQUEST) @app.post('/', include_in_schema=False) async def receive_message(request:fastapi.Request): ''' 接收微信公众号消息，并回复 ''' signature = request.query_params.get('msg_signature') timestamp = request.query_params.get('timestamp') nonce = request.query_params.get('nonce') if not signature or not timestamp or not nonce: logger.error('receive_message failed, missing parameters') return '', HTTPStatus.BAD_REQUEST logger.info('receive_message timestamp: %s, nonce: %s', timestamp, nonce) post_data = await request.body() code, post_data = get_wxcpt().DecryptMsg(post_data, signature, timestamp, nonce) if code != 0: logger.error('receive_message failed, error code: %s', code) return '', HTTPStatus.BAD_REQUEST xml = ET.fromstring(post_data) content = xml.find('Content').text from_user = xml.find('FromUserName').text asyncio.create_task(async_query_and_reply( apikey=config.bizwechat_config().qgi_api_key, model_name='qwen', database='contracts', question=content, to_user=from_user, )) logger.info('receive_message success, timestamp: %s, nonce: %s', timestamp, nonce) return '' from sqlcode.qgi import Executor, NewFormatter, ReturnType # from sqlcode.qwenapi import Generator from sqlcode.langchain_model import Generator from sqlcode.modelloader import ModelLoader, ModelManager async def async_query_and_reply(apikey:str, model_name:str, database:str, question:str, to_user:str): """ 根据请求执行查询并返回结果。 该函数接收一个模型名称、数据库名称和一个请求对象，通过这些信息来执行特定的查询操作。 查询的结果会根据请求中指定的返回类型进行处理和包装。 * :param model: 字符串类型，表示要使用的大模型的名称。例如 qwen-turbo 等 * :param database: 字符串类型，表示要查询的数据库的名称。 * :param req: Request 类型的对象，包含查询的具体问题和期望的返回类型。 """ logger.info('开始处理问题：%s', question) metadata = config.metadata(database) qwen_cfg = config.qwen_config("qwen_graph.conf") re_cfg = config.refineProblem_config() modelLoader = ModelLoader(config.model_config()) modelManager = ModelManager(modelLoader) modelManager.switch_model(model_name) model = modelManager.get_model() # db_dir = "chroma_db" # # 需要先检查向量数据库是否为最新，执行sqlcode/store_vecstore.py来更新向量数据库 # vectorstore = Chroma(persist_directory=db_dir, embedding_function=DashScopeEmbeddings()) # retriever = vectorstore.as_retriever(search_kwargs={'k': 2}) # re_prompt = PromptTemplate.from_template(re_cfg.prompt) # context_chain = retriever | format_docs # context = context_chain.invoke(question) # rag_chain = ( # {"context": lambda x: context, "question": RunnablePassthrough()} # | re_prompt # | model # | parse # ) # refine_question = rag_chain.invoke(question) # db = SQLDatabase.from_uri(metadata.connection_string) # toolkit = SQLDatabaseToolkit(db=db, llm=model) # remove_markdown_tool = StructuredTool.from_function( # func=remove_markdown_code_block, # description="当发生因为markdown标记导致的函数输入错误时，可以调用此函数来删除标记", # name="remove_markdown_code_block" # ) # prompt = PromptTemplate.from_template(qwen_cfg.prompt) # prompt = prompt.partial(metadata=metadata.metadata, product=metadata.product, example=qwen_cfg.params["example"], context=context) # agent_executor = create_sql_agent( # llm=model, # prompt=prompt, # toolkit=toolkit, # verbose=True, # agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, # extra_tools=[remove_markdown_tool], # agent_executor_kwargs={"handle_parsing_errors": True} # ) agent_executor = create_sql_graph(model=model, qwen_cfg=qwen_cfg, data_cfg=metadata) formatter = NewFormatter(format=ReturnType.WX_MD, tranlate=None, output_dir=config.osp.join(config.BASE_DIR, 'output'), site_url=config.app_config().site_url + '/output', ) generator = Generator(agentExcutor=agent_executor, messages=[{'role': 'system', 'content': qwen_cfg.system}], apikey=config.api_key(apikey).dashscope_api_key, seed=0, ) executor = Executor(generator=generator, formatter=formatter) # input = {"input": question, "refined_question": refine_question} input = {"messages": [("human", question)], "iterations": 0} ret = executor.query(connection_string=metadata.connection_string, input=input, ) # thought and result 分开发送因为消息大小限制为 2048 字节(utf-8) await send_msg(to_user, question, ret.error) if ret.thought.startswith('```') and ret.thought.endswith('```'): await send_msg(to_user, question, '针对这个问题，采用 SQL 查询：\n' + ret.thought) else: await send_msg(to_user, question, ret.thought) await send_msg(to_user, question, ret.result) async def send_msg(to_user:str, question:str, content:str): ''' 发送问题的回复给指定用户 ''' if not content: return msg = { "touser" : to_user, "msgtype": "markdown", "agentid" : config.bizwechat_config().agent_id, "markdown": {"content": content,}, } access_token = config.wxbiz_token() url = f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={access_token}" async with aiohttp.ClientSession() as session: async with session.post(url, json=msg) as resp: resp_text = await resp.text() logger.debug('send_msg to %s %s: %s', to_user, question, resp_text) if __name__ == '__main__': import sys if len(sys.argv) > 1 and sys.argv[1] == 'test': logging.basicConfig(level=logging.DEBUG) asyncio.run(async_query_and_reply(apikey=config.bizwechat_config().qgi_api_key, model='qwen-max', database='contracts', question="2020年业绩最好的分公司", to_user='SunHaiWen', )) elif len(sys.argv) > 1 and sys.argv[1] == 'send': logging.basicConfig(level=logging.DEBUG) asyncio.run(send_msg(to_user='SunHaiWen', question="2020年业绩最好的分公司", content="2020年业绩最好的分公司是北京分公司", )) else: import uvicorn uvicorn.run(app, host='0.0.0.0', port=9000)