refactor(prompts): 引用提示词增强

rag/prompts.py

@@ -32,7 +32,8 @@ def chunks_format(reference):
     def get_value(d, k1, k2):
         return d.get(k1, d.get(k2))
 
-    return [{
+    return [
+        {
             "id": get_value(chunk, "chunk_id", "id"),
             "content": get_value(chunk, "content", "content_with_weight"),
             "document_id": get_value(chunk, "doc_id", "document_id"),
@@ -40,8 +41,10 @@ def chunks_format(reference):
             "dataset_id": get_value(chunk, "kb_id", "dataset_id"),
             "image_id": get_value(chunk, "image_id", "img_id"),
             "positions": get_value(chunk, "positions", "position_int"),
-        "url": chunk.get("url")
-    } for chunk in reference.get("chunks", [])]
+            "url": chunk.get("url"),
+        }
+        for chunk in reference.get("chunks", [])
+    ]
 
 
 def llm_id2llm_type(llm_id):
@@ -65,13 +68,13 @@ def message_fit_in(msg, max_length=4000):
     返回:
         tuple: (实际token数, 调整后的消息列表)
     """
+
     def count():
         """计算当前消息列表的总token数"""
         nonlocal msg
         tks_cnts = []
         for m in msg:
-            tks_cnts.append(
-                {"role": m["role"], "count": num_tokens_from_string(m["content"])})
+            tks_cnts.append({"role": m["role"], "count": num_tokens_from_string(m["content"])})
         total = 0
         for m in tks_cnts:
             total += m["count"]
@@ -163,6 +166,10 @@ def citation_prompt():
 - 以格式 '##i$$ ##j$$'插入引用，其中 i, j 是所引用内容的 ID，并用 '##' 和 '$$' 包裹。
 - 在句子末尾插入引用，每个句子最多 4 个引用。
 - 如果答案内容不来自检索到的文本块，则不要插入引用。
+- 不要使用独立的文档 ID（例如 `#ID#`）。 
+- 在任何情况下，不得使用其他引用样式或格式（例如 `~~i==`、`[i]`、`(i)` 等）。 
+- 引用必须始终使用 `##i$$` 格式。
+- 任何未能遵守上述规则的情况，包括但不限于格式错误、使用禁止的样式或不支持的引用，都将被视为错误，应跳过为该句添加引用。
 
 --- 示例 ---
 <SYSTEM>: 以下是知识库:
@@ -210,10 +217,7 @@ def keyword_extraction(chat_mdl, content, topn=3):
 ### 文本内容
 {content}
 """
-    msg = [
-        {"role": "system", "content": prompt},
-        {"role": "user", "content": "Output: "}
-    ]
+    msg = [{"role": "system", "content": prompt}, {"role": "user", "content": "Output: "}]
     _, msg = message_fit_in(msg, chat_mdl.max_length)
     kwd = chat_mdl.chat(prompt, msg[1:], {"temperature": 0.2})
     if isinstance(kwd, tuple):
@@ -240,10 +244,7 @@ Requirements:
 {content}
 
 """
-    msg = [
-        {"role": "system", "content": prompt},
-        {"role": "user", "content": "Output: "}
-    ]
+    msg = [{"role": "system", "content": prompt}, {"role": "user", "content": "Output: "}]
     _, msg = message_fit_in(msg, chat_mdl.max_length)
     kwd = chat_mdl.chat(prompt, msg[1:], {"temperature": 0.2})
     if isinstance(kwd, tuple):
@@ -368,10 +369,7 @@ Output:
 {content}
 
 """
-    msg = [
-        {"role": "system", "content": prompt},
-        {"role": "user", "content": "Output: "}
-    ]
+    msg = [{"role": "system", "content": prompt}, {"role": "user", "content": "Output: "}]
     _, msg = message_fit_in(msg, chat_mdl.max_length)
     kwd = chat_mdl.chat(prompt, msg[1:], {"temperature": 0.5})
     if isinstance(kwd, tuple):
@@ -384,8 +382,8 @@ Output:
         return json_repair.loads(kwd)
     except json_repair.JSONDecodeError:
         try:
-            result = kwd.replace(prompt[:-1], '').replace('user', '').replace('model', '').strip()
-            result = '{' + result.split('{')[1].split('}')[0] + '}'
+            result = kwd.replace(prompt[:-1], "").replace("user", "").replace("model", "").strip()
+            result = "{" + result.split("{")[1].split("}")[0] + "}"
             return json_repair.loads(result)
         except Exception as e:
             logging.exception(f"JSON parsing error: {result} -> {e}")