Update app.py

app.py

@@ -1,10 +1,12 @@
 import gradio as gr
 from openai import OpenAI
 import os
-from fpdf import FPDF  # For PDF conversion
-from docx import Document  # For DOCX conversion
-import tempfile
+from io import BytesIO
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from docx import Document
 
+# Custom CSS
 css = '''
 .gradio-container{max-width: 1000px !important}
 h1{text-align:center}
@@ -13,6 +15,7 @@ footer {
 }
 '''
 
+# Set up OpenAI client
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
 
 client = OpenAI(
@@ -20,12 +23,17 @@ client = OpenAI(
     api_key=ACCESS_TOKEN,
 )
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    if history is None:
-        history = []
-
+# Function to handle chat responses
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
     messages = [{"role": "system", "content": system_message}]
-    
+
     for val in history:
         if val[0]:
             messages.append({"role": "user", "content": val[0]})
@@ -45,81 +53,120 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         messages=messages,
     ):
         token = message.choices[0].delta.content
-        
         response += token
         yield response
-    
-    history.append((message, response))
-    return history
 
-def save_as_file(history, conversion_type):
-    input_text = "\n".join([f"User: {h[0]}" for h in history if h[0]])
-    output_text = "\n".join([f"Assistant: {h[1]}" for h in history if h[1]])
-    file_name = None
+# Function to save chat history to a text file
+def save_as_txt(history):
+    with open("chat_history.txt", "w") as f:
+        for user_message, assistant_message in history:
+            f.write(f"User: {user_message}\n")
+            f.write(f"Assistant: {assistant_message}\n")
+    return "chat_history.txt"
+
+# Function to save chat history to a DOCX file
+def save_as_docx(history):
+    doc = Document()
+    doc.add_heading('Chat History', 0)
     
-    if conversion_type == "PDF":
-        pdf = FPDF()
-        pdf.add_page()
-        pdf.set_font("Arial", size=12)
-        pdf.multi_cell(0, 10, f"User Query: {input_text}\n\nResponse: {output_text}")
-        file_name = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf").name
-        pdf.output(file_name)
-    elif conversion_type == "DOCX":
-        doc = Document()
-        doc.add_heading('Conversation', 0)
-        doc.add_paragraph(f"User Query: {input_text}\n\nResponse: {output_text}")
-        file_name = tempfile.NamedTemporaryFile(delete=False, suffix=".docx").name
-        doc.save(file_name)
-    elif conversion_type == "TXT":
-        file_name = tempfile.NamedTemporaryFile(delete=False, suffix=".txt").name
-        with open(file_name, "w") as f:
-            f.write(f"User Query: {input_text}\n\nResponse: {output_text}")
+    for user_message, assistant_message in history:
+        doc.add_paragraph(f"User: {user_message}")
+        doc.add_paragraph(f"Assistant: {assistant_message}")
     
-    return file_name
-
-def convert_and_download(history, conversion_type):
-    if not history:
-        return None
-
-    file_path = save_as_file(history, conversion_type)
-    return file_path
-
-demo = gr.Blocks(css=css)
-
-with demo:
-    history_state = gr.State([])  # Initialize an empty list to store the conversation history
-
-    with gr.Row():
-        system_message = gr.Textbox(value="", label="System message")
-        max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
-        temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
-        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
+    doc.save("chat_history.docx")
+    return "chat_history.docx"
+
+# Function to save chat history to a PDF file
+def save_as_pdf(history):
+    buffer = BytesIO()
+    c = canvas.Canvas(buffer, pagesize=letter)
+    width, height = letter
+    y = height - 40
     
-    with gr.Row():
-        conversation_input = gr.Textbox(label="Message")
-        chat_output = gr.Textbox(label="Response", interactive=False)
+    c.drawString(30, y, "Chat History")
+    y -= 30
     
-    chat = gr.Button("Send")
-    conversion_type = gr.Dropdown(choices=["PDF", "DOCX", "TXT"], value="PDF", label="Conversion Type")
-    download_button = gr.Button("Convert and Download")
+    for user_message, assistant_message in history:
+        c.drawString(30, y, f"User: {user_message}")
+        y -= 20
+        c.drawString(30, y, f"Assistant: {assistant_message}")
+        y -= 30
+        
+        if y < 40:
+            c.showPage()
+            y = height - 40
+            
+    c.save()
+    buffer.seek(0)
     
-    file_output = gr.File()
+    with open("chat_history.pdf", "wb") as f:
+        f.write(buffer.read())
     
-    def update_history(message, history, system_message, max_tokens, temperature, top_p):
-        history = respond(message, history, system_message, max_tokens, temperature, top_p)
-        return history
-
-    chat.click(
-        fn=update_history,
-        inputs=[conversation_input, history_state, system_message, max_tokens, temperature, top_p],
-        outputs=[chat_output, history_state]
-    )
+    return "chat_history.pdf"
+
+# Gradio interface
+def handle_file_save(history, file_format):
+    if file_format == "txt":
+        return save_as_txt(history)
+    elif file_format == "docx":
+        return save_as_docx(history)
+    elif file_format == "pdf":
+        return save_as_pdf(history)
+
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-P",
+        ),
+        gr.Dropdown(
+            choices=["txt", "docx", "pdf"],
+            label="Save as",
+        ),
+    ],
+    outputs=[
+        gr.File(label="Download Chat History"),
+    ],
+    css=css,
+    theme="allenai/gradio-theme",
+)
 
-    download_button.click(
-        fn=convert_and_download,
-        inputs=[history_state, conversion_type],
-        outputs=file_output
-    )
+def save_handler(message, history, system_message, max_tokens, temperature, top_p, file_format):
+    response = respond(message, history, system_message, max_tokens, temperature, top_p)
+    saved_file = handle_file_save(history, file_format)
+    return saved_file
+
+demo = gr.Interface(
+    fn=save_handler,
+    inputs=[
+        gr.Textbox(value="", label="Message"),
+        gr.State(),
+        gr.Textbox(value="", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-P",
+        ),
+        gr.Dropdown(
+            choices=["txt", "docx", "pdf"],
+            label="Save as",
+        ),
+    ],
+    outputs=gr.File(label="Download Chat History"),
+    css=css,
+    theme="allenai/gradio-theme",
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()