diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py
index 607194db5..79fcfcad3 100644
--- a/api/db/services/dialog_service.py
+++ b/api/db/services/dialog_service.py
@@ -302,7 +302,7 @@ def chat(dialog, messages, stream=True, **kwargs):
     if "max_tokens" in gen_conf:
         gen_conf["max_tokens"] = min(gen_conf["max_tokens"], max_tokens - used_token_count)
 
-    def repair_bad_citation_formats(answer: str, kbinfos: dict, idx: dict):
+    def repair_bad_citation_formats(answer: str, kbinfos: dict, idx: set):
         max_index = len(kbinfos["chunks"])
 
         def safe_add(i):
@@ -623,4 +623,3 @@ def ask(question, kb_ids, tenant_id):
         answer = ans
         yield {"answer": answer, "reference": {}}
     yield decorate_answer(answer)
-