Wrap citation message in EventEmitter wrapper function.

Fixes issue introduced by PR #23.
Updates issue #8.

Author: EtiennePerot

open-webui/functions/run_code.py

@@ -488,6 +488,16 @@ async def message(self, content):
                 },
             )
 
+        async def citation(self, document, metadata, source):
+            await self._emit(
+                "citation",
+                {
+                    "document": document,
+                    "metadata": metadata,
+                    "source": source,
+                },
+            )
+
         async def code_execution_result(self, output):
             await self._emit(
                 "code_execution_result",
@@ -3280,7 +3290,8 @@ def _print_output(obj):
                 success = False
                 _print_output(e)
                 print(
-                    f"\u274c Self-test {name} failed: process failed: {e}", file=sys.stderr
+                    f"\u274c Self-test {name} failed: process failed: {e}",
+                    file=sys.stderr,
                 )
             except Exception as e:
                 success = False

open-webui/tools/run_code.py

@@ -220,6 +220,16 @@ async def message(self, content):
                 },
             )
 
+        async def citation(self, document, metadata, source):
+            await self._emit(
+                "citation",
+                {
+                    "document": document,
+                    "metadata": metadata,
+                    "source": source,
+                },
+            )
+
         async def code_execution_result(self, output):
             await self._emit(
                 "code_execution_result",
@@ -321,16 +331,8 @@ async def _fail(error_message, status="SANDBOX_ERROR"):
                     f"Running {language_title} code in gVisor sandbox..."
                 )
 
-                # TODO: Wrap this in an emitter helper function.
-                await event_emitter(
-                    {
-                        "type": "citation",
-                        "data": {
-                            "document": [code],
-                            "metadata": [code],
-                            "source": {"name": "run_code"},
-                        },
-                    }
+                await emitter.citation(
+                    document=[code], metadata=[code], source={"name": "run_code"}
                 )
 
                 try:
@@ -2516,12 +2518,15 @@ def _print_output(obj):
         except subprocess.CalledProcessError as e:
             success = False
             _print_output(e)
-            print(f"\u274c Self-test {name} failed: process failed: {e}", file=sys.stderr)
+            print(
+                f"\u274c Self-test {name} failed: process failed: {e}", file=sys.stderr
+            )
         except Exception as e:
             success = False
             exception_class = e.__class__
             print(
-                f"\u274c Self-test {name} failed: {exception_class}: {e}", file=sys.stderr
+                f"\u274c Self-test {name} failed: {exception_class}: {e}",
+                file=sys.stderr,
             )
         else:
             try: