92 lines
4.0 KiB
Python
92 lines
4.0 KiB
Python
|
from api.db import LLMType, ParserType
|
|||
|
|
from api.db.services.knowledgebase_service import KnowledgebaseService
|
|||
|
|
from api.db.services.llm_service import LLMBundle
|
|||
|
|
from api import settings
|
|||
|
|
from rag.app.tag import label_question
|
|||
|
|
from rag.prompts import kb_prompt
|
|||
|
|
|
|||
|
|
|
|||
|
|
def write_dialog(question, kb_ids, tenant_id):
|
|||
|
|
"""
|
|||
|
|
处理用户搜索请求，从知识库中检索相关信息并生成回答
|
|||
|
|
|
|||
|
|
参数:
|
|||
|
|
question (str): 用户的问题或查询
|
|||
|
|
kb_ids (list): 知识库ID列表，指定要搜索的知识库
|
|||
|
|
tenant_id (str): 租户ID，用于权限控制和资源隔离
|
|||
|
|
|
|||
|
|
流程:
|
|||
|
|
1. 获取指定知识库的信息
|
|||
|
|
2. 确定使用的嵌入模型
|
|||
|
|
3. 根据知识库类型选择检索器(普通检索器或知识图谱检索器)
|
|||
|
|
4. 初始化嵌入模型和聊天模型
|
|||
|
|
5. 执行检索操作获取相关文档片段
|
|||
|
|
6. 格式化知识库内容作为上下文
|
|||
|
|
7. 构建系统提示词
|
|||
|
|
8. 生成回答并添加引用标记
|
|||
|
|
9. 流式返回生成的回答
|
|||
|
|
|
|||
|
|
返回:
|
|||
|
|
generator: 生成器对象，产生包含回答和引用信息的字典
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
kbs = KnowledgebaseService.get_by_ids(kb_ids)
|
|||
|
|
embedding_list = list(set([kb.embd_id for kb in kbs]))
|
|||
|
|
|
|||
|
|
is_knowledge_graph = all([kb.parser_id == ParserType.KG for kb in kbs])
|
|||
|
|
retriever = settings.retrievaler if not is_knowledge_graph else settings.kg_retrievaler
|
|||
|
|
# 初始化嵌入模型，用于将文本转换为向量表示
|
|||
|
|
embd_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING, embedding_list[0])
|
|||
|
|
# 初始化聊天模型，用于生成回答
|
|||
|
|
chat_mdl = LLMBundle(tenant_id, LLMType.CHAT)
|
|||
|
|
# 获取聊天模型的最大token长度，用于控制上下文长度
|
|||
|
|
max_tokens = chat_mdl.max_length
|
|||
|
|
# 获取所有知识库的租户ID并去重
|
|||
|
|
tenant_ids = list(set([kb.tenant_id for kb in kbs]))
|
|||
|
|
# 调用检索器检索相关文档片段
|
|||
|
|
kbinfos = retriever.retrieval(question, embd_mdl, tenant_ids, kb_ids, 1, 12, 0.1, 0.3, aggs=False, rank_feature=label_question(question, kbs))
|
|||
|
|
# 将检索结果格式化为提示词，并确保不超过模型最大token限制
|
|||
|
|
knowledges = kb_prompt(kbinfos, max_tokens)
|
|||
|
|
prompt = """
|
|||
|
|
角色：你是一个聪明的助手。
|
|||
|
|
任务：总结知识库中的信息并回答用户的问题。
|
|||
|
|
要求与限制：
|
|||
|
|
- 绝不要捏造内容，尤其是数字。
|
|||
|
|
- 如果知识库中的信息与用户问题无关，**只需回答：对不起，未提供相关信息。
|
|||
|
|
- 使用Markdown格式进行回答。
|
|||
|
|
- 使用用户提问所用的语言作答。
|
|||
|
|
- 绝不要捏造内容，尤其是数字。
|
|||
|
|
|
|||
|
|
### 来自知识库的信息
|
|||
|
|
%s
|
|||
|
|
|
|||
|
|
以上是来自知识库的信息。
|
|||
|
|
|
|||
|
|
""" % "\n".join(knowledges)
|
|||
|
|
msg = [{"role": "user", "content": question}]
|
|||
|
|
|
|||
|
|
# 生成完成后添加回答中的引用标记
|
|||
|
|
# def decorate_answer(answer):
|
|||
|
|
# nonlocal knowledges, kbinfos, prompt
|
|||
|
|
# answer, idx = retriever.insert_citations(answer, [ck["content_ltks"] for ck in kbinfos["chunks"]], [ck["vector"] for ck in kbinfos["chunks"]], embd_mdl, tkweight=0.7, vtweight=0.3)
|
|||
|
|
# idx = set([kbinfos["chunks"][int(i)]["doc_id"] for i in idx])
|
|||
|
|
# recall_docs = [d for d in kbinfos["doc_aggs"] if d["doc_id"] in idx]
|
|||
|
|
# if not recall_docs:
|
|||
|
|
# recall_docs = kbinfos["doc_aggs"]
|
|||
|
|
# kbinfos["doc_aggs"] = recall_docs
|
|||
|
|
# refs = deepcopy(kbinfos)
|
|||
|
|
# for c in refs["chunks"]:
|
|||
|
|
# if c.get("vector"):
|
|||
|
|
# del c["vector"]
|
|||
|
|
|
|||
|
|
# if answer.lower().find("invalid key") >= 0 or answer.lower().find("invalid api") >= 0:
|
|||
|
|
# answer += " Please set LLM API-Key in 'User Setting -> Model Providers -> API-Key'"
|
|||
|
|
# refs["chunks"] = chunks_format(refs)
|
|||
|
|
# return {"answer": answer, "reference": refs}
|
|||
|
|
|
|||
|
|
answer = ""
|
|||
|
|
for ans in chat_mdl.chat_streamly(prompt, msg, {"temperature": 0.1}):
|
|||
|
|
answer = ans
|
|||
|
|
yield {"answer": answer, "reference": {}}
|
|||
|
|
# yield decorate_answer(answer)
|