134 lines
6.7 KiB
Python
134 lines
6.7 KiB
Python
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from datetime import datetime
|
||
import json
|
||
from uuid import uuid4
|
||
from fastapi import Body, HTTPException
|
||
from fastapi.responses import FileResponse
|
||
from typing import Any, Dict, List, Optional
|
||
from sse_starlette.sse import EventSourceResponse
|
||
from langchain.chains import LLMChain
|
||
from configs.comparison_config import comparison
|
||
from configs import LLM_MODELS
|
||
from configs.model_config import MAX_TOKENS
|
||
from server.agent.tools.draw_plot import create_and_save_plot
|
||
from server.callback_handler.conversation_callback_handler import ConversationCallbackHandler
|
||
from server.chat import utils
|
||
from server.chat.policy_fun_iast import get_llm_model_response
|
||
from typing import AsyncIterable
|
||
from configs.basic_config import *
|
||
from server.chat.utils import History
|
||
from server.custom.AsyncIteratorCallbackHandlerNew import AsyncIteratorCallbackHandler
|
||
from server.utils import get_ChatOpenAI, get_prompt_template, wrap_done
|
||
from langchain.prompts.chat import ChatPromptTemplate
|
||
|
||
async def chat_comparison_test(
|
||
uid: Optional[str] = Body(None, description="时间"),
|
||
content_list: Optional[List[str]] = Body(None, description="对比文献内容，需要是个包含所有文件内容的字符串列表"),
|
||
):
|
||
async def chat_iterator() -> AsyncIterable[str]:
|
||
nonlocal content_list
|
||
|
||
# 按照对比模板提取的对比维度
|
||
context = comparison[0]["summary"]
|
||
|
||
# 用于处理提炼资料的进度条
|
||
total_tasks = len(content_list)
|
||
completed_tasks = 0
|
||
|
||
|
||
# 用于存储从原文资料里提炼出来的所有有用的信息
|
||
useful_content = []
|
||
|
||
# 使用线程池执行异步任务对文献所有内容进行提取
|
||
with ThreadPoolExecutor() as executor:
|
||
# 创建一个任务列表，每个任务处理 content_list 中的一个元素
|
||
futures = [
|
||
executor.submit(
|
||
get_llm_model_response,
|
||
strategy_name="query rewrite",
|
||
llm_model_name=LLM_MODELS[1],
|
||
template_prompt_name="extract_key_points",
|
||
prompt_param_dict={"time": datetime.now().strftime("%Y%m%d"), "context": context, "content": content},
|
||
temperature=0.01,
|
||
max_tokens=512
|
||
)
|
||
for content in content_list
|
||
]
|
||
|
||
# 收集所有任务的结果
|
||
for future in as_completed(futures):
|
||
useful_content.append(future.result())
|
||
completed_tasks += 1
|
||
progress = (completed_tasks / total_tasks) * 50
|
||
print(f"Progress: {progress:.2f}%")
|
||
yield json.dumps(
|
||
{"progress": f"{progress:.2f}%"},
|
||
ensure_ascii=False)
|
||
|
||
# 对所有结果进行处理
|
||
useful_res = "\n".join(f"第{i+1}篇文章:\n{content}\n"for i,content in enumerate(useful_content) )
|
||
|
||
# 根据要点数据先制作出对比图
|
||
pic_param=get_llm_model_response(
|
||
strategy_name="query rewrite",
|
||
llm_model_name=LLM_MODELS[0],
|
||
template_prompt_name="make_comparison_pic",
|
||
prompt_param_dict={"time": datetime.now().strftime("%Y%m%d"), "context": useful_res, "content": "【参数说明】:参数格式：以JSON格式提供， 必需参数：\n 1. data：图表数据格式如下{\"XXX\": XX, \"XXX\": XX, \"XXX\": X, \"XXX\": X}其中XXX代表分类名称，xx代表分类数据量。\n2.title：图表标题3.xlabel：横轴标题（你按照分的几类属于哪个大类）必须有\n4.ylabel：纵轴标题（你的数值数据是什么）必须有\n 5.plot_type:图表类型必须从以下几个当中选一个作为输入【bar,pie,line】其中bar代表柱状图pie代表饼图line代表折线图你只能选一个作为输入，\n【使用指南】：\n当需要生成图表时，请使用此工具。\n输入必须包含以下参数：:{\n\"data\": {\"XXX\": XX, \"XXX\": XX, \"XXX\": X, \"XXX\": X},\"title\": \"X\",\"xlabel\": \"X\",\"ylabel\": \"X\",\"plot_type\": \"X\"}\n示例：{'\"data\": {\"Category A\": 23, \"Category B\": 17, \"Category C\": 35, \"Category D\": 29},\"title\": \"My Chart\",\"xlabel\": \"Category\",\"ylabel\": \"Value\",\"plot_type\": \"pie\"}\n请务必以json格式输入方便使用。\n" },
|
||
temperature=0.01,
|
||
max_tokens=MAX_TOKENS
|
||
).replace("```json","").replace("```","").strip()
|
||
pic_params = json.loads(pic_param)
|
||
pic_content={}
|
||
uuid = str(uuid4())
|
||
utils.set_shared_variable(uuid, pic_content)
|
||
temp = {}
|
||
temp["uuid"] = uuid
|
||
pic_res = ""
|
||
for param in pic_params:
|
||
pic_res+=create_and_save_plot(f"{param}{temp}")+"\n"
|
||
#pic_res = create_and_save_plot(f"{pic_param}{temp}")
|
||
# 生成对比文件
|
||
callback = AsyncIteratorCallbackHandler()
|
||
callbacks = [callback]
|
||
conversation_callback = ConversationCallbackHandler(conversation_id="", message_id="",
|
||
chat_type="llm_chat",
|
||
query=useful_res)
|
||
callbacks.append(conversation_callback)
|
||
model = get_ChatOpenAI(
|
||
model_name=LLM_MODELS[0],
|
||
temperature=0.01,
|
||
max_tokens=MAX_TOKENS,
|
||
callbacks=callbacks,
|
||
)
|
||
memory = None
|
||
prompt_template = get_prompt_template("comparison_chat", "write_report")
|
||
input_prompt = History(role="system", content=prompt_template).to_msg_template(False)
|
||
chat_prompt = ChatPromptTemplate.from_messages([input_prompt])
|
||
chain = LLMChain(prompt=chat_prompt, llm=model)
|
||
task = asyncio.create_task(wrap_done(
|
||
chain.acall({
|
||
# "context": res,
|
||
"context": comparison[0]["content"],
|
||
"time": datetime.now().strftime("%Y%m%d"),
|
||
"content": useful_res,
|
||
"pic": pic_res,
|
||
}),
|
||
callback.done),
|
||
)
|
||
history_summary = ""
|
||
# if stream:
|
||
async for token in callback.aiter():
|
||
# Use server-sent-events to stream the response
|
||
history_summary += token
|
||
# yield json.dumps(
|
||
# {"text": token},
|
||
# ensure_ascii=False)
|
||
yield json.dumps(
|
||
{"text": history_summary},
|
||
ensure_ascii=False)
|
||
yield json.dumps(
|
||
{"progress": "100%"},
|
||
ensure_ascii=False)
|
||
await task
|
||
|
||
return EventSourceResponse(chat_iterator()) |