import uuid from fastapi import Body from langchain.memory import ( CombinedMemory, ConversationBufferMemory, ConversationSummaryMemory, ConversationBufferWindowMemory ) from configs.model_config import MAX_TOKENS from configs.outline_config import outlines from langchain.chains.question_answering import load_qa_chain from typing import Any from sse_starlette.sse import EventSourceResponse from configs import LLM_MODELS, TEMPERATURE, HISTORY_LEN from configs.kb_config import CH_BASE_NAME from server.agent.tools.search_tool import search_tool from server.chat import utils from server.chat.agent_chat_test import agent_chat_test from server.chat.policy_fun_iast import get_llm_model_response from server.chat.solve_problem import solve_problem from server.knowledge_base.kb_service.base import TextRank from server.knowledge_base.model.kb_document_model import DocumentWithVSId from server.utils import wrap_done, get_ChatOpenAI from langchain.chains import LLMChain, ConversationChain from langchain.callbacks import AsyncIteratorCallbackHandler from typing import AsyncIterable import asyncio import json from langchain.prompts.chat import ChatPromptTemplate from typing import List, Optional, Union from server.chat.utils import History, compute_lps, remove_after_and_including,remove_before_and_including from langchain.prompts import PromptTemplate from server.utils import get_prompt_template, get_format_template from server.memory.conversation_db_buffer_memory import ConversationBufferDBMemory from server.db.repository import add_message_to_db from server.callback_handler.conversation_callback_handler import ConversationCallbackHandler from datetime import datetime from langchain_core.messages import SystemMessage import time as t from configs.basic_config import * async def process_task(task): results = [] async for result in task: results.append(result) return results async def knowledge_chat_test( uid: Optional[Any] = Body(None, description="用户ID"), rag: bool = Body(False, description="增强标识, 当knowledge_base_list有值时，需要传True"), knowledge_base_list: Optional[List[str]] = Body(None, description="知识库名称列表"), query: str = Body(..., description="用户输入", examples=["恼羞成怒"]), conversation_id: str = Body("", description="对话框ID"), model_name: str = Body(LLM_MODELS[0], description="LLM 模型名称。"), temperature: float = Body(TEMPERATURE, description="LLM 采样温度", ge=0.0, le=2.0), max_tokens: Optional[int] = Body( MAX_TOKENS, description="限制LLM生成Token数量，默认None代表模型最大值" ), # top_p: float = Body(TOP_P, description="LLM 核采样。勿与temperature 同时设置", gt=0.0, lt=1.0), prompt_name: str = Body("default", description="使用的prompt模板名称(在configs/prompt_config.py中配置)"), ): async def chat_iterator() -> AsyncIterable[str]: nonlocal rag,max_tokens, uid, model_name,knowledge_base_list # rag = True max_tokens = MAX_TOKENS history_summary = "" memory = None message_id = str(uuid.uuid1())+"q" think_type = {"text": "", "message_id": message_id} user_prompt_name = prompt_name outline_num = "" res = "" callback = AsyncIteratorCallbackHandler() callbacks = [callback] conversation_callback = ConversationCallbackHandler(conversation_id=conversation_id, message_id=message_id, chat_type="llm_chat", query=query) callbacks.append(conversation_callback) time_based_uuid = str(uuid.uuid1())+"q" tip ={} tip["END"]="" tip["source_docs"]=[] tip["num"] = 0 tip["title"]=[] utils.set_shared_variable(time_based_uuid,tip) # 先判断需不需要使用知识库或者联网检索同时需要获取大纲模板 task1 = solve_problem(user_prompt_name= user_prompt_name,query=query, conversation_id="", history=[], model_name=model_name,temperature=temperature,max_tokens=max_tokens,prompt_name="think_route",stream=True) task2 = solve_problem(user_prompt_name= user_prompt_name,query=query, conversation_id="", history=[], model_name=model_name,temperature=temperature,max_tokens=max_tokens,prompt_name="outlines_route",stream=True) results = await asyncio.gather(process_task(task1), process_task(task2)) for result0 in results[0]: think_type["text"] += json.loads(result0)["text"].strip() break for result1 in results[1]: outline_num += json.loads(result1)["text"] break if(str(think_type["text"]) == "2") or rag == True: yield json.dumps({"text": "思考中..."}, ensure_ascii=False) model = get_ChatOpenAI( model_name=model_name, temperature=temperature, max_tokens=max_tokens, callbacks=callbacks, ) search_query = get_llm_model_response( strategy_name="query rewrite", llm_model_name=LLM_MODELS[0], template_prompt_name="query_rewrite", prompt_param_dict={"query": query, "history": [], "time": datetime.now().strftime("%Y%m%d")}, temperature=0.01, max_tokens=512 ) # 第一个 JSON 对象 keys = json.loads(search_query).keys() keys_list = list(keys) if rag: knowledge_base_list = list(set(CH_BASE_NAME).union(knowledge_base_list)) logger.info(f"大纲撰写参考的个人知识库名称：{knowledge_base_list}") first_json = { "query": json.loads(search_query)[keys_list[0]], "knowledge_name": knowledge_base_list, "keywords": [] } else: first_json = { "query": json.loads(search_query)[keys_list[0]], "knowledge_name": CH_BASE_NAME, "keywords": [] } # 第二个 JSON 对象 second_json = { "uuid": time_based_uuid } origin_query=json.dumps(first_json) + json.dumps(second_json) # 截断检索资料，避免超token res += search_tool(origin_query) if len(res) >= 7000: res = res[:7000] else: model = get_ChatOpenAI( model_name=LLM_MODELS[0], temperature=temperature, max_tokens=max_tokens, callbacks=callbacks, ) prompt_template = get_prompt_template("knowledge_base_chat", "default_outlines") input_prompt = History(role="system", content=prompt_template).to_msg_template(False) # input_msg = History(role="user", content=query).to_msg_template(False) chat_prompt = ChatPromptTemplate.from_messages([input_prompt]) chain = LLMChain(prompt=chat_prompt, llm=model) input_prompt = History(role="system", content=prompt_template).to_msg_template(False) chat_prompt = ChatPromptTemplate.from_messages([input_prompt]) format_template = get_format_template("knowledge_base_chat", "default_outlines") chain = load_qa_chain( model, chain_type="stuff", memory=memory, prompt=chat_prompt, verbose=True ) # docs = list(itertools.chain(policydocs, reportdocs, journaldocs, personaldocs)) if not outline_num or outline_num == "" or "无" in outline_num: outline_num = len(outlines)-1 if res =="": res+="根据你的自身能力回答" results1=[] results1.append(DocumentWithVSId(page_content=res)) logger.info(f"outline_num:{outline_num}") task = asyncio.create_task(wrap_done( chain.acall({ # "context": res, "input_documents": results1, "self_knowledge":outlines[int(outline_num)], "history": [], "question": query, "format_template": format_template, "time": datetime.now().strftime("%Y%m%d") }), callback.done), ) # if stream: async for token in callback.aiter(): # Use server-sent-events to stream the response history_summary += token yield json.dumps( {"text": token, "message_id": message_id}, ensure_ascii=False) # try: # source = utils.get_shared_variable(time_based_uuid) # source_docs = source["source_docs"] # docs_string = "\n"+"\n".join(f"{str(doc)}\n" for doc in source_docs) # yield json.dumps({"docs": docs_string}, ensure_ascii=False) # except Exception as e: # logging.error(f"获取知识库联想信息失败：{e}") # print("知识库联想存在异常或没有进行知识库联想") # pass #summary = TextRank(history_summary, 80) #yield json.dumps({"summary":summary}, ensure_ascii=False) # del question_history await task callback = AsyncIteratorCallbackHandler() return EventSourceResponse(chat_iterator())