From 913e38dffa50abd0687f8d4495b89412811e67dd Mon Sep 17 00:00:00 2001
From: renxin <90580890+renxinx@users.noreply.github.com>
Date: Thu, 3 Apr 2025 21:20:21 -0700
Subject: [PATCH] Feature/revise docs ci (#5056)

---
 docs/wrap_run_llm.py | 37 ++++++++++++++++++++++---------------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/docs/wrap_run_llm.py b/docs/wrap_run_llm.py
index 0d0fe8b25..2d21442cf 100644
--- a/docs/wrap_run_llm.py
+++ b/docs/wrap_run_llm.py
@@ -13,27 +13,34 @@ def insert_runllm_widget(html_content):
     return re.sub(r"</body>", f"{widget_script}\n</body>", html_content)
 
 
+def process_html_files(build_dir):
+    for root, dirs, files in os.walk(build_dir):
+        for file in files:
+            if file.endswith(".html"):
+                file_path = os.path.join(root, file)
+
+                # Read the HTML file
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+
+                # Insert the RunLLM widget
+                modified_content = insert_runllm_widget(content)
+
+                # Write back the modified content
+                with open(file_path, "w", encoding="utf-8") as f:
+                    f.write(modified_content)
+
+
 def main():
     # Get the build directory path
     build_dir = os.path.join(
         os.path.dirname(os.path.abspath(__file__)), "_build", "html"
     )
-    index_file = os.path.join(build_dir, "index.html")
-
-    # Process only index.html
-    if os.path.exists(index_file):
-        # Read the HTML file
-        with open(index_file, "r", encoding="utf-8") as f:
-            content = f.read()
-
-        # Insert the RunLLM widget
-        modified_content = insert_runllm_widget(content)
-
-        # Write back the modified content
-        with open(index_file, "w", encoding="utf-8") as f:
-            f.write(modified_content)
+    # Process all HTML files
+    if os.path.exists(build_dir):
+        process_html_files(build_dir)
     else:
-        print(f"Index file not found: {index_file}")
+        print(f"Build directory not found: {build_dir}")
 
 
 if __name__ == "__main__":