From 8b1e2ce2795b753778ca5d6408a05fc7ed0adcbb Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Mon, 10 Jun 2024 17:12:48 -0700
Subject: [PATCH] feat: code format

---
 backend/apps/webui/routers/utils.py           | 17 ++++++++
 src/lib/apis/utils/index.ts                   | 33 +++++++++++++++
 src/lib/components/common/CodeEditor.svelte   | 42 ++++++++++++++++++-
 .../workspace/Tools/CodeEditor.svelte         |  2 +
 4 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/backend/apps/webui/routers/utils.py b/backend/apps/webui/routers/utils.py
index 18491130a..8f6d663b4 100644
--- a/backend/apps/webui/routers/utils.py
+++ b/backend/apps/webui/routers/utils.py
@@ -7,6 +7,8 @@ from pydantic import BaseModel
 
 from fpdf import FPDF
 import markdown
+import black
+
 
 from apps.webui.internal.db import DB
 from utils.utils import get_admin_user
@@ -26,6 +28,21 @@ async def get_gravatar(
     return get_gravatar_url(email)
 
 
+class CodeFormatRequest(BaseModel):
+    code: str
+
+
+@router.post("/code/format")
+async def format_code(request: CodeFormatRequest):
+    try:
+        formatted_code = black.format_str(request.code, mode=black.Mode())
+        return {"code": formatted_code}
+    except black.NothingChanged:
+        return {"code": request.code}
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
 class MarkdownForm(BaseModel):
     md: str
 
diff --git a/src/lib/apis/utils/index.ts b/src/lib/apis/utils/index.ts
index acb4a04aa..f99ac5641 100644
--- a/src/lib/apis/utils/index.ts
+++ b/src/lib/apis/utils/index.ts
@@ -22,6 +22,39 @@ export const getGravatarUrl = async (email: string) => {
 	return res;
 };
 
+export const formatPythonCode = async (code: string) => {
+	let error = null;
+
+	const res = await fetch(`${WEBUI_API_BASE_URL}/utils/code/format`, {
+		method: 'POST',
+		headers: {
+			'Content-Type': 'application/json'
+		},
+		body: JSON.stringify({
+			code: code
+		})
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.log(err);
+
+			error = err;
+			if (err.detail) {
+				error = err.detail;
+			}
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
 export const downloadChatAsPDF = async (chat: object) => {
 	let error = null;
 
diff --git a/src/lib/components/common/CodeEditor.svelte b/src/lib/components/common/CodeEditor.svelte
index f25b3c736..5cf10e1e4 100644
--- a/src/lib/components/common/CodeEditor.svelte
+++ b/src/lib/components/common/CodeEditor.svelte
@@ -6,10 +6,15 @@
 	import { acceptCompletion } from '@codemirror/autocomplete';
 	import { indentWithTab } from '@codemirror/commands';
 
+	import { indentUnit } from '@codemirror/language';
 	import { python } from '@codemirror/lang-python';
 	import { oneDark } from '@codemirror/theme-one-dark';
 
-	import { onMount } from 'svelte';
+	import { onMount, createEventDispatcher } from 'svelte';
+	import { formatPythonCode } from '$lib/apis/utils';
+	import { toast } from 'svelte-sonner';
+
+	const dispatch = createEventDispatcher();
 
 	export let boilerplate = '';
 	export let value = '';
@@ -19,10 +24,31 @@
 	let isDarkMode = false;
 	let editorTheme = new Compartment();
 
+	const formatPythonCodeHandler = async () => {
+		if (codeEditor) {
+			console.log('formatPythonCodeHandler');
+			const res = await formatPythonCode(value).catch((error) => {
+				toast.error(error);
+				return null;
+			});
+
+			if (res && res.code) {
+				const formattedCode = res.code;
+				codeEditor.dispatch({
+					changes: [{ from: 0, to: codeEditor.state.doc.length, insert: formattedCode }]
+				});
+				return true;
+			}
+
+			return false;
+		}
+	};
+
 	let extensions = [
 		basicSetup,
 		keymap.of([{ key: 'Tab', run: acceptCompletion }, indentWithTab]),
 		python(),
+		indentUnit.of('    '),
 		placeholder('Enter your code here...'),
 		EditorView.updateListener.of((e) => {
 			if (e.docChanged) {
@@ -78,8 +104,22 @@
 			attributeFilter: ['class']
 		});
 
+		// Add a keyboard shortcut to format the code when Ctrl/Cmd + S is pressed
+		// Override the default browser save functionality
+
+		const handleSave = (e) => {
+			if ((e.ctrlKey || e.metaKey) && e.key === 's') {
+				e.preventDefault();
+				formatPythonCodeHandler();
+				dispatch('save');
+			}
+		};
+
+		document.addEventListener('keydown', handleSave);
+
 		return () => {
 			observer.disconnect();
+			document.removeEventListener('keydown', handleSave);
 		};
 	});
 </script>
diff --git a/src/lib/components/workspace/Tools/CodeEditor.svelte b/src/lib/components/workspace/Tools/CodeEditor.svelte
index 1696e4bf8..ec7656096 100644
--- a/src/lib/components/workspace/Tools/CodeEditor.svelte
+++ b/src/lib/components/workspace/Tools/CodeEditor.svelte
@@ -7,6 +7,8 @@
 # Use Sphinx-style docstrings to document your tools, they will be used for generating tools specifications
 # Please refer to function_calling_filter_pipeline.py file from pipelines project for an example
 
+# Tip: Use Ctrl/Cmd + S to format the code
+
 from datetime import datetime
 import requests