您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

dialog_service.py 30KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728
  1. #
  2. # Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. #
  16. import binascii
  17. import logging
  18. import re
  19. import time
  20. from copy import deepcopy
  21. from datetime import datetime
  22. from functools import partial
  23. from timeit import default_timer as timer
  24. from langfuse import Langfuse
  25. from peewee import fn
  26. from agentic_reasoning import DeepResearcher
  27. from api import settings
  28. from api.db import LLMType, ParserType, StatusEnum
  29. from api.db.db_models import DB, Dialog
  30. from api.db.services.common_service import CommonService
  31. from api.db.services.document_service import DocumentService
  32. from api.db.services.knowledgebase_service import KnowledgebaseService
  33. from api.db.services.langfuse_service import TenantLangfuseService
  34. from api.db.services.llm_service import LLMBundle, TenantLLMService
  35. from api.utils import current_timestamp, datetime_format
  36. from rag.app.resume import forbidden_select_fields4resume
  37. from rag.app.tag import label_question
  38. from rag.nlp.search import index_name
  39. from rag.prompts import chunks_format, citation_prompt, cross_languages, full_question, kb_prompt, keyword_extraction, message_fit_in
  40. from rag.prompts.prompts import gen_meta_filter
  41. from rag.utils import num_tokens_from_string, rmSpace
  42. from rag.utils.tavily_conn import Tavily
  43. class DialogService(CommonService):
  44. model = Dialog
  45. @classmethod
  46. def save(cls, **kwargs):
  47. """Save a new record to database.
  48. This method creates a new record in the database with the provided field values,
  49. forcing an insert operation rather than an update.
  50. Args:
  51. **kwargs: Record field values as keyword arguments.
  52. Returns:
  53. Model instance: The created record object.
  54. """
  55. sample_obj = cls.model(**kwargs).save(force_insert=True)
  56. return sample_obj
  57. @classmethod
  58. def update_many_by_id(cls, data_list):
  59. """Update multiple records by their IDs.
  60. This method updates multiple records in the database, identified by their IDs.
  61. It automatically updates the update_time and update_date fields for each record.
  62. Args:
  63. data_list (list): List of dictionaries containing record data to update.
  64. Each dictionary must include an 'id' field.
  65. """
  66. with DB.atomic():
  67. for data in data_list:
  68. data["update_time"] = current_timestamp()
  69. data["update_date"] = datetime_format(datetime.now())
  70. cls.model.update(data).where(cls.model.id == data["id"]).execute()
  71. @classmethod
  72. @DB.connection_context()
  73. def get_list(cls, tenant_id, page_number, items_per_page, orderby, desc, id, name):
  74. chats = cls.model.select()
  75. if id:
  76. chats = chats.where(cls.model.id == id)
  77. if name:
  78. chats = chats.where(cls.model.name == name)
  79. chats = chats.where((cls.model.tenant_id == tenant_id) & (cls.model.status == StatusEnum.VALID.value))
  80. if desc:
  81. chats = chats.order_by(cls.model.getter_by(orderby).desc())
  82. else:
  83. chats = chats.order_by(cls.model.getter_by(orderby).asc())
  84. chats = chats.paginate(page_number, items_per_page)
  85. return list(chats.dicts())
  86. @classmethod
  87. @DB.connection_context()
  88. def get_by_tenant_ids(cls, joined_tenant_ids, user_id, page_number, items_per_page, orderby, desc, keywords, parser_id=None):
  89. from api.db.db_models import User
  90. fields = [
  91. cls.model.id,
  92. cls.model.tenant_id,
  93. cls.model.name,
  94. cls.model.description,
  95. cls.model.language,
  96. cls.model.llm_id,
  97. cls.model.llm_setting,
  98. cls.model.prompt_type,
  99. cls.model.prompt_config,
  100. cls.model.similarity_threshold,
  101. cls.model.vector_similarity_weight,
  102. cls.model.top_n,
  103. cls.model.top_k,
  104. cls.model.do_refer,
  105. cls.model.rerank_id,
  106. cls.model.kb_ids,
  107. cls.model.status,
  108. User.nickname,
  109. User.avatar.alias("tenant_avatar"),
  110. cls.model.update_time,
  111. cls.model.create_time,
  112. ]
  113. if keywords:
  114. dialogs = (
  115. cls.model.select(*fields)
  116. .join(User, on=(cls.model.tenant_id == User.id))
  117. .where(
  118. (cls.model.tenant_id.in_(joined_tenant_ids) | (cls.model.tenant_id == user_id)) & (cls.model.status == StatusEnum.VALID.value),
  119. (fn.LOWER(cls.model.name).contains(keywords.lower())),
  120. )
  121. )
  122. else:
  123. dialogs = (
  124. cls.model.select(*fields)
  125. .join(User, on=(cls.model.tenant_id == User.id))
  126. .where(
  127. (cls.model.tenant_id.in_(joined_tenant_ids) | (cls.model.tenant_id == user_id)) & (cls.model.status == StatusEnum.VALID.value),
  128. )
  129. )
  130. if parser_id:
  131. dialogs = dialogs.where(cls.model.parser_id == parser_id)
  132. if desc:
  133. dialogs = dialogs.order_by(cls.model.getter_by(orderby).desc())
  134. else:
  135. dialogs = dialogs.order_by(cls.model.getter_by(orderby).asc())
  136. count = dialogs.count()
  137. if page_number and items_per_page:
  138. dialogs = dialogs.paginate(page_number, items_per_page)
  139. return list(dialogs.dicts()), count
  140. def chat_solo(dialog, messages, stream=True):
  141. if TenantLLMService.llm_id2llm_type(dialog.llm_id) == "image2text":
  142. chat_mdl = LLMBundle(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id)
  143. else:
  144. chat_mdl = LLMBundle(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
  145. prompt_config = dialog.prompt_config
  146. tts_mdl = None
  147. if prompt_config.get("tts"):
  148. tts_mdl = LLMBundle(dialog.tenant_id, LLMType.TTS)
  149. msg = [{"role": m["role"], "content": re.sub(r"##\d+\$\$", "", m["content"])} for m in messages if m["role"] != "system"]
  150. if stream:
  151. last_ans = ""
  152. delta_ans = ""
  153. for ans in chat_mdl.chat_streamly(prompt_config.get("system", ""), msg, dialog.llm_setting):
  154. answer = ans
  155. delta_ans = ans[len(last_ans) :]
  156. if num_tokens_from_string(delta_ans) < 16:
  157. continue
  158. last_ans = answer
  159. yield {"answer": answer, "reference": {}, "audio_binary": tts(tts_mdl, delta_ans), "prompt": "", "created_at": time.time()}
  160. delta_ans = ""
  161. if delta_ans:
  162. yield {"answer": answer, "reference": {}, "audio_binary": tts(tts_mdl, delta_ans), "prompt": "", "created_at": time.time()}
  163. else:
  164. answer = chat_mdl.chat(prompt_config.get("system", ""), msg, dialog.llm_setting)
  165. user_content = msg[-1].get("content", "[content not available]")
  166. logging.debug("User: {}|Assistant: {}".format(user_content, answer))
  167. yield {"answer": answer, "reference": {}, "audio_binary": tts(tts_mdl, answer), "prompt": "", "created_at": time.time()}
  168. def get_models(dialog):
  169. embd_mdl, chat_mdl, rerank_mdl, tts_mdl = None, None, None, None
  170. kbs = KnowledgebaseService.get_by_ids(dialog.kb_ids)
  171. embedding_list = list(set([kb.embd_id for kb in kbs]))
  172. if len(embedding_list) > 1:
  173. raise Exception("**ERROR**: Knowledge bases use different embedding models.")
  174. if embedding_list:
  175. embd_mdl = LLMBundle(dialog.tenant_id, LLMType.EMBEDDING, embedding_list[0])
  176. if not embd_mdl:
  177. raise LookupError("Embedding model(%s) not found" % embedding_list[0])
  178. if TenantLLMService.llm_id2llm_type(dialog.llm_id) == "image2text":
  179. chat_mdl = LLMBundle(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id)
  180. else:
  181. chat_mdl = LLMBundle(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
  182. if dialog.rerank_id:
  183. rerank_mdl = LLMBundle(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id)
  184. if dialog.prompt_config.get("tts"):
  185. tts_mdl = LLMBundle(dialog.tenant_id, LLMType.TTS)
  186. return kbs, embd_mdl, rerank_mdl, chat_mdl, tts_mdl
  187. BAD_CITATION_PATTERNS = [
  188. re.compile(r"\(\s*ID\s*[: ]*\s*(\d+)\s*\)"), # (ID: 12)
  189. re.compile(r"\[\s*ID\s*[: ]*\s*(\d+)\s*\]"), # [ID: 12]
  190. re.compile(r"【\s*ID\s*[: ]*\s*(\d+)\s*】"), # 【ID: 12】
  191. re.compile(r"ref\s*(\d+)", flags=re.IGNORECASE), # ref12、REF 12
  192. ]
  193. def repair_bad_citation_formats(answer: str, kbinfos: dict, idx: set):
  194. max_index = len(kbinfos["chunks"])
  195. def safe_add(i):
  196. if 0 <= i < max_index:
  197. idx.add(i)
  198. return True
  199. return False
  200. def find_and_replace(pattern, group_index=1, repl=lambda i: f"ID:{i}", flags=0):
  201. nonlocal answer
  202. def replacement(match):
  203. try:
  204. i = int(match.group(group_index))
  205. if safe_add(i):
  206. return f"[{repl(i)}]"
  207. except Exception:
  208. pass
  209. return match.group(0)
  210. answer = re.sub(pattern, replacement, answer, flags=flags)
  211. for pattern in BAD_CITATION_PATTERNS:
  212. find_and_replace(pattern)
  213. return answer, idx
  214. def meta_filter(metas: dict, filters: list[dict]):
  215. doc_ids = []
  216. def filter_out(v2docs, operator, value):
  217. nonlocal doc_ids
  218. for input,docids in v2docs.items():
  219. try:
  220. input = float(input)
  221. value = float(value)
  222. except Exception:
  223. input = str(input)
  224. value = str(value)
  225. for conds in [
  226. (operator == "contains", str(value).lower() in str(input).lower()),
  227. (operator == "not contains", str(value).lower() not in str(input).lower()),
  228. (operator == "start with", str(input).lower().startswith(str(value).lower())),
  229. (operator == "end with", str(input).lower().endswith(str(value).lower())),
  230. (operator == "empty", not input),
  231. (operator == "not empty", input),
  232. (operator == "=", input == value),
  233. (operator == "≠", input != value),
  234. (operator == ">", input > value),
  235. (operator == "<", input < value),
  236. (operator == "≥", input >= value),
  237. (operator == "≤", input <= value),
  238. ]:
  239. try:
  240. if all(conds):
  241. doc_ids.extend(docids)
  242. except Exception:
  243. pass
  244. for k, v2docs in metas.items():
  245. for f in filters:
  246. if k != f["key"]:
  247. continue
  248. filter_out(v2docs, f["op"], f["value"])
  249. return doc_ids
  250. def chat(dialog, messages, stream=True, **kwargs):
  251. assert messages[-1]["role"] == "user", "The last content of this conversation is not from user."
  252. if not dialog.kb_ids and not dialog.prompt_config.get("tavily_api_key"):
  253. for ans in chat_solo(dialog, messages, stream):
  254. yield ans
  255. return
  256. chat_start_ts = timer()
  257. if TenantLLMService.llm_id2llm_type(dialog.llm_id) == "image2text":
  258. llm_model_config = TenantLLMService.get_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id)
  259. else:
  260. llm_model_config = TenantLLMService.get_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
  261. max_tokens = llm_model_config.get("max_tokens", 8192)
  262. check_llm_ts = timer()
  263. langfuse_tracer = None
  264. trace_context = {}
  265. langfuse_keys = TenantLangfuseService.filter_by_tenant(tenant_id=dialog.tenant_id)
  266. if langfuse_keys:
  267. langfuse = Langfuse(public_key=langfuse_keys.public_key, secret_key=langfuse_keys.secret_key, host=langfuse_keys.host)
  268. if langfuse.auth_check():
  269. langfuse_tracer = langfuse
  270. trace_id = langfuse_tracer.create_trace_id()
  271. trace_context = {"trace_id": trace_id}
  272. check_langfuse_tracer_ts = timer()
  273. kbs, embd_mdl, rerank_mdl, chat_mdl, tts_mdl = get_models(dialog)
  274. toolcall_session, tools = kwargs.get("toolcall_session"), kwargs.get("tools")
  275. if toolcall_session and tools:
  276. chat_mdl.bind_tools(toolcall_session, tools)
  277. bind_models_ts = timer()
  278. retriever = settings.retrievaler
  279. questions = [m["content"] for m in messages if m["role"] == "user"][-3:]
  280. attachments = kwargs["doc_ids"].split(",") if "doc_ids" in kwargs else []
  281. if "doc_ids" in messages[-1]:
  282. attachments = messages[-1]["doc_ids"]
  283. prompt_config = dialog.prompt_config
  284. field_map = KnowledgebaseService.get_field_map(dialog.kb_ids)
  285. # try to use sql if field mapping is good to go
  286. if field_map:
  287. logging.debug("Use SQL to retrieval:{}".format(questions[-1]))
  288. ans = use_sql(questions[-1], field_map, dialog.tenant_id, chat_mdl, prompt_config.get("quote", True))
  289. if ans:
  290. yield ans
  291. return
  292. for p in prompt_config["parameters"]:
  293. if p["key"] == "knowledge":
  294. continue
  295. if p["key"] not in kwargs and not p["optional"]:
  296. raise KeyError("Miss parameter: " + p["key"])
  297. if p["key"] not in kwargs:
  298. prompt_config["system"] = prompt_config["system"].replace("{%s}" % p["key"], " ")
  299. if len(questions) > 1 and prompt_config.get("refine_multiturn"):
  300. questions = [full_question(dialog.tenant_id, dialog.llm_id, messages)]
  301. else:
  302. questions = questions[-1:]
  303. if prompt_config.get("cross_languages"):
  304. questions = [cross_languages(dialog.tenant_id, dialog.llm_id, questions[0], prompt_config["cross_languages"])]
  305. if dialog.meta_data_filter:
  306. metas = DocumentService.get_meta_by_kbs(dialog.kb_ids)
  307. if dialog.meta_data_filter.get("method") == "auto":
  308. filters = gen_meta_filter(chat_mdl, metas, questions[-1])
  309. attachments.extend(meta_filter(metas, filters))
  310. elif dialog.meta_data_filter.get("method") == "manual":
  311. attachments.extend(meta_filter(metas, dialog.meta_data_filter["manual"]))
  312. if prompt_config.get("keyword", False):
  313. questions[-1] += keyword_extraction(chat_mdl, questions[-1])
  314. refine_question_ts = timer()
  315. thought = ""
  316. kbinfos = {"total": 0, "chunks": [], "doc_aggs": []}
  317. if "knowledge" not in [p["key"] for p in prompt_config["parameters"]]:
  318. knowledges = []
  319. else:
  320. tenant_ids = list(set([kb.tenant_id for kb in kbs]))
  321. knowledges = []
  322. if prompt_config.get("reasoning", False):
  323. reasoner = DeepResearcher(
  324. chat_mdl,
  325. prompt_config,
  326. partial(retriever.retrieval, embd_mdl=embd_mdl, tenant_ids=tenant_ids, kb_ids=dialog.kb_ids, page=1, page_size=dialog.top_n, similarity_threshold=0.2, vector_similarity_weight=0.3),
  327. )
  328. for think in reasoner.thinking(kbinfos, " ".join(questions)):
  329. if isinstance(think, str):
  330. thought = think
  331. knowledges = [t for t in think.split("\n") if t]
  332. elif stream:
  333. yield think
  334. else:
  335. if embd_mdl:
  336. kbinfos = retriever.retrieval(
  337. " ".join(questions),
  338. embd_mdl,
  339. tenant_ids,
  340. dialog.kb_ids,
  341. 1,
  342. dialog.top_n,
  343. dialog.similarity_threshold,
  344. dialog.vector_similarity_weight,
  345. doc_ids=attachments,
  346. top=dialog.top_k,
  347. aggs=False,
  348. rerank_mdl=rerank_mdl,
  349. rank_feature=label_question(" ".join(questions), kbs),
  350. )
  351. if prompt_config.get("tavily_api_key"):
  352. tav = Tavily(prompt_config["tavily_api_key"])
  353. tav_res = tav.retrieve_chunks(" ".join(questions))
  354. kbinfos["chunks"].extend(tav_res["chunks"])
  355. kbinfos["doc_aggs"].extend(tav_res["doc_aggs"])
  356. if prompt_config.get("use_kg"):
  357. ck = settings.kg_retrievaler.retrieval(" ".join(questions), tenant_ids, dialog.kb_ids, embd_mdl, LLMBundle(dialog.tenant_id, LLMType.CHAT))
  358. if ck["content_with_weight"]:
  359. kbinfos["chunks"].insert(0, ck)
  360. knowledges = kb_prompt(kbinfos, max_tokens)
  361. logging.debug("{}->{}".format(" ".join(questions), "\n->".join(knowledges)))
  362. retrieval_ts = timer()
  363. if not knowledges and prompt_config.get("empty_response"):
  364. empty_res = prompt_config["empty_response"]
  365. yield {"answer": empty_res, "reference": kbinfos, "prompt": "\n\n### Query:\n%s" % " ".join(questions), "audio_binary": tts(tts_mdl, empty_res)}
  366. return {"answer": prompt_config["empty_response"], "reference": kbinfos}
  367. kwargs["knowledge"] = "\n------\n" + "\n\n------\n\n".join(knowledges)
  368. gen_conf = dialog.llm_setting
  369. msg = [{"role": "system", "content": prompt_config["system"].format(**kwargs)}]
  370. prompt4citation = ""
  371. if knowledges and (prompt_config.get("quote", True) and kwargs.get("quote", True)):
  372. prompt4citation = citation_prompt()
  373. msg.extend([{"role": m["role"], "content": re.sub(r"##\d+\$\$", "", m["content"])} for m in messages if m["role"] != "system"])
  374. used_token_count, msg = message_fit_in(msg, int(max_tokens * 0.95))
  375. assert len(msg) >= 2, f"message_fit_in has bug: {msg}"
  376. prompt = msg[0]["content"]
  377. if "max_tokens" in gen_conf:
  378. gen_conf["max_tokens"] = min(gen_conf["max_tokens"], max_tokens - used_token_count)
  379. def decorate_answer(answer):
  380. nonlocal embd_mdl, prompt_config, knowledges, kwargs, kbinfos, prompt, retrieval_ts, questions, langfuse_tracer
  381. refs = []
  382. ans = answer.split("</think>")
  383. think = ""
  384. if len(ans) == 2:
  385. think = ans[0] + "</think>"
  386. answer = ans[1]
  387. if knowledges and (prompt_config.get("quote", True) and kwargs.get("quote", True)):
  388. idx = set([])
  389. if embd_mdl and not re.search(r"\[ID:([0-9]+)\]", answer):
  390. answer, idx = retriever.insert_citations(
  391. answer,
  392. [ck["content_ltks"] for ck in kbinfos["chunks"]],
  393. [ck["vector"] for ck in kbinfos["chunks"]],
  394. embd_mdl,
  395. tkweight=1 - dialog.vector_similarity_weight,
  396. vtweight=dialog.vector_similarity_weight,
  397. )
  398. else:
  399. for match in re.finditer(r"\[ID:([0-9]+)\]", answer):
  400. i = int(match.group(1))
  401. if i < len(kbinfos["chunks"]):
  402. idx.add(i)
  403. answer, idx = repair_bad_citation_formats(answer, kbinfos, idx)
  404. idx = set([kbinfos["chunks"][int(i)]["doc_id"] for i in idx])
  405. recall_docs = [d for d in kbinfos["doc_aggs"] if d["doc_id"] in idx]
  406. if not recall_docs:
  407. recall_docs = kbinfos["doc_aggs"]
  408. kbinfos["doc_aggs"] = recall_docs
  409. refs = deepcopy(kbinfos)
  410. for c in refs["chunks"]:
  411. if c.get("vector"):
  412. del c["vector"]
  413. if answer.lower().find("invalid key") >= 0 or answer.lower().find("invalid api") >= 0:
  414. answer += " Please set LLM API-Key in 'User Setting -> Model providers -> API-Key'"
  415. finish_chat_ts = timer()
  416. total_time_cost = (finish_chat_ts - chat_start_ts) * 1000
  417. check_llm_time_cost = (check_llm_ts - chat_start_ts) * 1000
  418. check_langfuse_tracer_cost = (check_langfuse_tracer_ts - check_llm_ts) * 1000
  419. bind_embedding_time_cost = (bind_models_ts - check_langfuse_tracer_ts) * 1000
  420. refine_question_time_cost = (refine_question_ts - bind_models_ts) * 1000
  421. retrieval_time_cost = (retrieval_ts - refine_question_ts) * 1000
  422. generate_result_time_cost = (finish_chat_ts - retrieval_ts) * 1000
  423. tk_num = num_tokens_from_string(think + answer)
  424. prompt += "\n\n### Query:\n%s" % " ".join(questions)
  425. prompt = (
  426. f"{prompt}\n\n"
  427. "## Time elapsed:\n"
  428. f" - Total: {total_time_cost:.1f}ms\n"
  429. f" - Check LLM: {check_llm_time_cost:.1f}ms\n"
  430. f" - Check Langfuse tracer: {check_langfuse_tracer_cost:.1f}ms\n"
  431. f" - Bind models: {bind_embedding_time_cost:.1f}ms\n"
  432. f" - Query refinement(LLM): {refine_question_time_cost:.1f}ms\n"
  433. f" - Retrieval: {retrieval_time_cost:.1f}ms\n"
  434. f" - Generate answer: {generate_result_time_cost:.1f}ms\n\n"
  435. "## Token usage:\n"
  436. f" - Generated tokens(approximately): {tk_num}\n"
  437. f" - Token speed: {int(tk_num / (generate_result_time_cost / 1000.0))}/s"
  438. )
  439. # Add a condition check to call the end method only if langfuse_tracer exists
  440. if langfuse_tracer and "langfuse_generation" in locals():
  441. langfuse_output = "\n" + re.sub(r"^.*?(### Query:.*)", r"\1", prompt, flags=re.DOTALL)
  442. langfuse_output = {"time_elapsed:": re.sub(r"\n", " \n", langfuse_output), "created_at": time.time()}
  443. langfuse_generation.update(output=langfuse_output)
  444. langfuse_generation.end()
  445. return {"answer": think + answer, "reference": refs, "prompt": re.sub(r"\n", " \n", prompt), "created_at": time.time()}
  446. if langfuse_tracer:
  447. langfuse_generation = langfuse_tracer.start_generation(
  448. trace_context=trace_context, name="chat", model=llm_model_config["llm_name"], input={"prompt": prompt, "prompt4citation": prompt4citation, "messages": msg}
  449. )
  450. if stream:
  451. last_ans = ""
  452. answer = ""
  453. for ans in chat_mdl.chat_streamly(prompt + prompt4citation, msg[1:], gen_conf):
  454. if thought:
  455. ans = re.sub(r"^.*</think>", "", ans, flags=re.DOTALL)
  456. answer = ans
  457. delta_ans = ans[len(last_ans) :]
  458. if num_tokens_from_string(delta_ans) < 16:
  459. continue
  460. last_ans = answer
  461. yield {"answer": thought + answer, "reference": {}, "audio_binary": tts(tts_mdl, delta_ans)}
  462. delta_ans = answer[len(last_ans) :]
  463. if delta_ans:
  464. yield {"answer": thought + answer, "reference": {}, "audio_binary": tts(tts_mdl, delta_ans)}
  465. yield decorate_answer(thought + answer)
  466. else:
  467. answer = chat_mdl.chat(prompt + prompt4citation, msg[1:], gen_conf)
  468. user_content = msg[-1].get("content", "[content not available]")
  469. logging.debug("User: {}|Assistant: {}".format(user_content, answer))
  470. res = decorate_answer(answer)
  471. res["audio_binary"] = tts(tts_mdl, answer)
  472. yield res
  473. def use_sql(question, field_map, tenant_id, chat_mdl, quota=True):
  474. sys_prompt = "You are a Database Administrator. You need to check the fields of the following tables based on the user's list of questions and write the SQL corresponding to the last question."
  475. user_prompt = """
  476. Table name: {};
  477. Table of database fields are as follows:
  478. {}
  479. Question are as follows:
  480. {}
  481. Please write the SQL, only SQL, without any other explanations or text.
  482. """.format(index_name(tenant_id), "\n".join([f"{k}: {v}" for k, v in field_map.items()]), question)
  483. tried_times = 0
  484. def get_table():
  485. nonlocal sys_prompt, user_prompt, question, tried_times
  486. sql = chat_mdl.chat(sys_prompt, [{"role": "user", "content": user_prompt}], {"temperature": 0.06})
  487. sql = re.sub(r"^.*</think>", "", sql, flags=re.DOTALL)
  488. logging.debug(f"{question} ==> {user_prompt} get SQL: {sql}")
  489. sql = re.sub(r"[\r\n]+", " ", sql.lower())
  490. sql = re.sub(r".*select ", "select ", sql.lower())
  491. sql = re.sub(r" +", " ", sql)
  492. sql = re.sub(r"([;;]|```).*", "", sql)
  493. if sql[: len("select ")] != "select ":
  494. return None, None
  495. if not re.search(r"((sum|avg|max|min)\(|group by )", sql.lower()):
  496. if sql[: len("select *")] != "select *":
  497. sql = "select doc_id,docnm_kwd," + sql[6:]
  498. else:
  499. flds = []
  500. for k in field_map.keys():
  501. if k in forbidden_select_fields4resume:
  502. continue
  503. if len(flds) > 11:
  504. break
  505. flds.append(k)
  506. sql = "select doc_id,docnm_kwd," + ",".join(flds) + sql[8:]
  507. logging.debug(f"{question} get SQL(refined): {sql}")
  508. tried_times += 1
  509. return settings.retrievaler.sql_retrieval(sql, format="json"), sql
  510. tbl, sql = get_table()
  511. if tbl is None:
  512. return None
  513. if tbl.get("error") and tried_times <= 2:
  514. user_prompt = """
  515. Table name: {};
  516. Table of database fields are as follows:
  517. {}
  518. Question are as follows:
  519. {}
  520. Please write the SQL, only SQL, without any other explanations or text.
  521. The SQL error you provided last time is as follows:
  522. {}
  523. Error issued by database as follows:
  524. {}
  525. Please correct the error and write SQL again, only SQL, without any other explanations or text.
  526. """.format(index_name(tenant_id), "\n".join([f"{k}: {v}" for k, v in field_map.items()]), question, sql, tbl["error"])
  527. tbl, sql = get_table()
  528. logging.debug("TRY it again: {}".format(sql))
  529. logging.debug("GET table: {}".format(tbl))
  530. if tbl.get("error") or len(tbl["rows"]) == 0:
  531. return None
  532. docid_idx = set([ii for ii, c in enumerate(tbl["columns"]) if c["name"] == "doc_id"])
  533. doc_name_idx = set([ii for ii, c in enumerate(tbl["columns"]) if c["name"] == "docnm_kwd"])
  534. column_idx = [ii for ii in range(len(tbl["columns"])) if ii not in (docid_idx | doc_name_idx)]
  535. # compose Markdown table
  536. columns = (
  537. "|" + "|".join([re.sub(r"(/.*|([^()]+))", "", field_map.get(tbl["columns"][i]["name"], tbl["columns"][i]["name"])) for i in column_idx]) + ("|Source|" if docid_idx and docid_idx else "|")
  538. )
  539. line = "|" + "|".join(["------" for _ in range(len(column_idx))]) + ("|------|" if docid_idx and docid_idx else "")
  540. rows = ["|" + "|".join([rmSpace(str(r[i])) for i in column_idx]).replace("None", " ") + "|" for r in tbl["rows"]]
  541. rows = [r for r in rows if re.sub(r"[ |]+", "", r)]
  542. if quota:
  543. rows = "\n".join([r + f" ##{ii}$$ |" for ii, r in enumerate(rows)])
  544. else:
  545. rows = "\n".join([r + f" ##{ii}$$ |" for ii, r in enumerate(rows)])
  546. rows = re.sub(r"T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+Z)?\|", "|", rows)
  547. if not docid_idx or not doc_name_idx:
  548. logging.warning("SQL missing field: " + sql)
  549. return {"answer": "\n".join([columns, line, rows]), "reference": {"chunks": [], "doc_aggs": []}, "prompt": sys_prompt}
  550. docid_idx = list(docid_idx)[0]
  551. doc_name_idx = list(doc_name_idx)[0]
  552. doc_aggs = {}
  553. for r in tbl["rows"]:
  554. if r[docid_idx] not in doc_aggs:
  555. doc_aggs[r[docid_idx]] = {"doc_name": r[doc_name_idx], "count": 0}
  556. doc_aggs[r[docid_idx]]["count"] += 1
  557. return {
  558. "answer": "\n".join([columns, line, rows]),
  559. "reference": {
  560. "chunks": [{"doc_id": r[docid_idx], "docnm_kwd": r[doc_name_idx]} for r in tbl["rows"]],
  561. "doc_aggs": [{"doc_id": did, "doc_name": d["doc_name"], "count": d["count"]} for did, d in doc_aggs.items()],
  562. },
  563. "prompt": sys_prompt,
  564. }
  565. def tts(tts_mdl, text):
  566. if not tts_mdl or not text:
  567. return
  568. bin = b""
  569. for chunk in tts_mdl.tts(text):
  570. bin += chunk
  571. return binascii.hexlify(bin).decode("utf-8")
  572. def ask(question, kb_ids, tenant_id, chat_llm_name=None):
  573. kbs = KnowledgebaseService.get_by_ids(kb_ids)
  574. embedding_list = list(set([kb.embd_id for kb in kbs]))
  575. is_knowledge_graph = all([kb.parser_id == ParserType.KG for kb in kbs])
  576. retriever = settings.retrievaler if not is_knowledge_graph else settings.kg_retrievaler
  577. embd_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING, embedding_list[0])
  578. chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, chat_llm_name)
  579. max_tokens = chat_mdl.max_length
  580. tenant_ids = list(set([kb.tenant_id for kb in kbs]))
  581. kbinfos = retriever.retrieval(question, embd_mdl, tenant_ids, kb_ids, 1, 12, 0.1, 0.3, aggs=False, rank_feature=label_question(question, kbs))
  582. knowledges = kb_prompt(kbinfos, max_tokens)
  583. prompt = """
  584. Role: You're a smart assistant. Your name is Miss R.
  585. Task: Summarize the information from knowledge bases and answer user's question.
  586. Requirements and restriction:
  587. - DO NOT make things up, especially for numbers.
  588. - If the information from knowledge is irrelevant with user's question, JUST SAY: Sorry, no relevant information provided.
  589. - Answer with markdown format text.
  590. - Answer in language of user's question.
  591. - DO NOT make things up, especially for numbers.
  592. ### Information from knowledge bases
  593. %s
  594. The above is information from knowledge bases.
  595. """ % "\n".join(knowledges)
  596. msg = [{"role": "user", "content": question}]
  597. def decorate_answer(answer):
  598. nonlocal knowledges, kbinfos, prompt
  599. 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)
  600. idx = set([kbinfos["chunks"][int(i)]["doc_id"] for i in idx])
  601. recall_docs = [d for d in kbinfos["doc_aggs"] if d["doc_id"] in idx]
  602. if not recall_docs:
  603. recall_docs = kbinfos["doc_aggs"]
  604. kbinfos["doc_aggs"] = recall_docs
  605. refs = deepcopy(kbinfos)
  606. for c in refs["chunks"]:
  607. if c.get("vector"):
  608. del c["vector"]
  609. if answer.lower().find("invalid key") >= 0 or answer.lower().find("invalid api") >= 0:
  610. answer += " Please set LLM API-Key in 'User Setting -> Model Providers -> API-Key'"
  611. refs["chunks"] = chunks_format(refs)
  612. return {"answer": answer, "reference": refs}
  613. answer = ""
  614. for ans in chat_mdl.chat_streamly(prompt, msg, {"temperature": 0.1}):
  615. answer = ans
  616. yield {"answer": answer, "reference": {}}
  617. yield decorate_answer(answer)