download data for different formats

app.py

@@ -4,6 +4,8 @@ from utils import (extract_wiki_id, get_wiki_details,
                    get_translate_prompt)
 import json
 import json_repair
+import os
+import tempfile
 
 # Define language options for translation
 LANGUAGES = {
@@ -162,6 +164,73 @@ def format_debug_info(debug_info):
     
     return markdown
 
+# Functions to generate downloadable content for original and translated articles
+def generate_download_original(download_format, article_title, article_summary, article_xml, sections):
+    """
+    Generate a downloadable original content file in the specified format.
+    """
+    # Prepare content and filename
+    if download_format == "Wikipedia XML":
+        content = article_xml or ""
+        filename = f"{article_title or 'article'}.xml"
+    elif download_format == "HTML":
+        parts = [f"<h1>{article_title}</h1>", f"<p>{article_summary}</p>"]
+        for title, text in sections.items():
+            parts.append(f"<h2>{title}</h2>")
+            parts.append(f"<p>{text}</p>")
+        content = "\n".join(parts)
+        filename = f"{article_title or 'article'}.html"
+    elif download_format == "JSON":
+        obj = {"title": article_title, "summary": article_summary, "sections": sections}
+        content = json.dumps(obj, ensure_ascii=False, indent=2)
+        filename = f"{article_title or 'article'}.json"
+    else:  # Plain Text
+        parts = [article_title or 'Article', article_summary or '']
+        for title, text in sections.items():
+            parts.append(f"## {title}\n{text}")
+        content = "\n\n".join(parts)
+        filename = f"{article_title or 'article'}.txt"
+    # Write to a temp file and return its path
+    temp_path = os.path.join(tempfile.gettempdir(), filename)
+    with open(temp_path, 'w', encoding='utf-8') as f:
+        f.write(content)
+    return temp_path
+
+def generate_download_translated(download_format, article_title, article_summary, translations):
+    """
+    Generate a downloadable translated content file in the specified format.
+    """
+    # Prepare content and filename
+    if download_format == "Wikipedia XML":
+        parts = [f"<article title=\"{article_title}\">"]
+        for title, text in translations.items():
+            parts.append(f"  <section title=\"{title}\">{text}</section>")
+        parts.append("</article>")
+        content = "\n".join(parts)
+        filename = f"{article_title or 'article'}_translated.xml"
+    elif download_format == "HTML":
+        parts = [f"<h1>{article_title}</h1>", f"<p>{article_summary}</p>"]
+        for title, text in translations.items():
+            parts.append(f"<h2>{title}</h2>")
+            parts.append(f"<p>{text}</p>")
+        content = "\n".join(parts)
+        filename = f"{article_title or 'article'}_translated.html"
+    elif download_format == "JSON":
+        obj = {"title": article_title, "summary": article_summary, "sections": translations}
+        content = json.dumps(obj, ensure_ascii=False, indent=2)
+        filename = f"{article_title or 'article'}_translated.json"
+    else:  # Plain Text
+        parts = [article_title or 'Article', article_summary or '']
+        for title, text in translations.items():
+            parts.append(f"## {title}\n{text}")
+        content = "\n\n".join(parts)
+        filename = f"{article_title or 'article'}_translated.txt"
+    # Write to a temp file and return its path
+    temp_path = os.path.join(tempfile.gettempdir(), filename)
+    with open(temp_path, 'w', encoding='utf-8') as f:
+        f.write(content)
+    return temp_path
+
 # Add this function to update UI with sections from Wikipedia content
 def update_ui_with_sections(sections):
     """
@@ -428,8 +497,18 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
         outputs=[sidebar_expanded, sidebar, main_content, sidebar_show_btn]
     )
     
-    # Add debug info display to the bottom of the sidebar
+    # Add download options to the bottom of the sidebar
     with sidebar:
+        download_format = gr.Dropdown(
+            choices=["Wikipedia XML", "HTML", "JSON", "Plain Text"],
+            value="Wikipedia XML",
+            label="Download Format"
+        )
+        download_original_btn = gr.Button("Download Original")
+        download_original_file = gr.File(label="Original Article")
+        download_translated_btn = gr.Button("Download Translated")
+        download_translated_file = gr.File(label="Translated Article")
+        # Debug info display
         debug_header = gr.Markdown("### Debug Information", visible=False)
         debug_display = gr.Markdown(visible=False)
     
@@ -446,6 +525,18 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
             fn=lambda: (gr.update(visible=True), gr.update(visible=True)),
             outputs=[debug_header, debug_display]
         )
+
+    # Connect download buttons
+    download_original_btn.click(
+        fn=generate_download_original,
+        inputs=[download_format, article_title, aticle_summary, article_xml, sections_state],
+        outputs=[download_original_file]
+    )
+    download_translated_btn.click(
+        fn=generate_download_translated,
+        inputs=[download_format, article_title, aticle_summary, sections_state],
+        outputs=[download_translated_file]
+    )
     
 # Launch the app
 if __name__ == "__main__":