From 21a28e3bf094976fe86ae605f2f3d5be2175fb08 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Fri, 25 Oct 2024 21:58:29 -0700
Subject: [PATCH] enh: more robust knowledge file info retrieval

---
 .../apps/webui/routers/knowledge.py           | 41 +++++++++++++++----
 1 file changed, 33 insertions(+), 8 deletions(-)

diff --git a/backend/open_webui/apps/webui/routers/knowledge.py b/backend/open_webui/apps/webui/routers/knowledge.py
index 9cb38a821..fe5d27e3f 100644
--- a/backend/open_webui/apps/webui/routers/knowledge.py
+++ b/backend/open_webui/apps/webui/routers/knowledge.py
@@ -47,15 +47,40 @@ async def get_knowledge_items(
                 detail=ERROR_MESSAGES.NOT_FOUND,
             )
     else:
-        return [
-            KnowledgeResponse(
-                **knowledge.model_dump(),
-                files=Files.get_file_metadatas_by_ids(
-                    knowledge.data.get("file_ids", []) if knowledge.data else []
-                ),
+        knowledge_bases = []
+
+        for knowledge in Knowledges.get_knowledge_items():
+            files = Files.get_file_metadatas_by_ids(
+                knowledge.data.get("file_ids", []) if knowledge.data else []
             )
-            for knowledge in Knowledges.get_knowledge_items()
-        ]
+
+            # Check if all files exist
+            if len(files) != len(knowledge.data.get("file_ids", [])):
+                missing_files = list(
+                    set(knowledge.data.get("file_ids", []))
+                    - set([file.id for file in files])
+                )
+                if missing_files:
+                    data = knowledge.data or {}
+                    file_ids = data.get("file_ids", [])
+
+                    for missing_file in missing_files:
+                        file_ids.remove(missing_file)
+
+                    data["file_ids"] = file_ids
+                    Knowledges.update_knowledge_by_id(
+                        id=knowledge.id, form_data=KnowledgeUpdateForm(data=data)
+                    )
+
+                    files = Files.get_file_metadatas_by_ids(file_ids)
+
+            knowledge_bases.append(
+                KnowledgeResponse(
+                    **knowledge.model_dump(),
+                    files=files,
+                )
+            )
+        return knowledge_bases
 
 
 ############################