feat: new frontend with model download function (#5008)

2025-05-18 20:00:17 +07:00 · 2025-05-18 20:00:17 +07:00 · c1091ce812
commit c1091ce812
parent 1b2a29565f
10 changed files with 318 additions and 71 deletions
--- a/web-app/src/containers/ChatInput.tsx
+++ b/web-app/src/containers/ChatInput.tsx
@ -41,6 +41,7 @@ import { listen } from '@tauri-apps/api/event'
 import { SystemEvent } from '@/types/events'
 import { CompletionMessagesBuilder } from '@/lib/messages'
 import { ChatCompletionMessageToolCall } from 'openai/resources'
+import { getTools } from '@/services/mcp'

 type ChatInputProps = {
  className?: string
@ -93,16 +94,15 @@ const ChatInput = ({ className, showSpeedToken = true }: ChatInputProps) => {
  }, [])

  useEffect(() => {
-    window.core?.api?.getTools().then((data: MCPTool[]) => {
-      setTools(data)
-    })
-
-    let unsubscribe = () => {}
-    listen(SystemEvent.MCP_UPDATE, () => {
-      window.core?.api?.getTools().then((data: MCPTool[]) => {
+    function updateTools() {
+      getTools().then((data: MCPTool[]) => {
        setTools(data)
      })
-    }).then((unsub) => {
+    }
+    updateTools()
+
+    let unsubscribe = () => {}
+    listen(SystemEvent.MCP_UPDATE, updateTools).then((unsub) => {
      // Unsubscribe from the event when the component unmounts
      unsubscribe = unsub
    })
@ -199,8 +199,11 @@ const ChatInput = ({ className, showSpeedToken = true }: ChatInputProps) => {
          accumulatedText
        )
        builder.addAssistantMessage(accumulatedText, undefined, toolCalls)
-        const updatedMessage = await postMessageProcessing(toolCalls, builder, finalContent)
-        console.log(updatedMessage)
+        const updatedMessage = await postMessageProcessing(
+          toolCalls,
+          builder,
+          finalContent
+        )
        addMessage(updatedMessage ?? finalContent)

        isCompleted = !toolCalls.length
--- a/web-app/src/containers/DownloadManegement.tsx
+++ b/web-app/src/containers/DownloadManegement.tsx
@ -4,25 +4,122 @@ import {
  PopoverTrigger,
 } from '@/components/ui/popover'
 import { Progress } from '@/components/ui/progress'
+import { useDownloadStore } from '@/hooks/useDownloadStore'
+import { abortDownload } from '@/services/models'
+import { DownloadEvent, DownloadState, events } from '@janhq/core'
 import { IconPlayerPauseFilled, IconX } from '@tabler/icons-react'
+import { useCallback, useEffect, useMemo } from 'react'

 export function DownloadManagement() {
+  const { downloads, updateProgress, removeDownload } = useDownloadStore()
+  const downloadCount = useMemo(
+    () => Object.keys(downloads).length,
+    [downloads]
+  )
+  const downloadProcesses = useMemo(
+    () =>
+      Object.values(downloads).map((download) => ({
+        id: download.id,
+        name: download.name,
+        progress: download.progress,
+        current: download.current,
+        total: download.total,
+      })),
+    [downloads]
+  )
+
+  const overallProgress = useMemo(() => {
+    const total = downloadProcesses.reduce((acc, download) => {
+      return acc + download.total
+    }, 0)
+    const current = downloadProcesses.reduce((acc, download) => {
+      return acc + download.current
+    }, 0)
+    return total > 0 ? current / total : 0
+  }, [downloadProcesses])
+
+  const onFileDownloadUpdate = useCallback(
+    async (state: DownloadState) => {
+      console.debug('onFileDownloadUpdate', state)
+      updateProgress(
+        state.modelId,
+        state.percent,
+        state.modelId,
+        state.size?.transferred,
+        state.size?.total
+      )
+    },
+    [updateProgress]
+  )
+
+  const onFileDownloadError = useCallback(
+    (state: DownloadState) => {
+      console.debug('onFileDownloadError', state)
+      removeDownload(state.modelId)
+    },
+    [removeDownload]
+  )
+
+  const onFileDownloadStopped = useCallback(
+    (state: DownloadState) => {
+      console.debug('onFileDownloadError', state)
+      removeDownload(state.modelId)
+    },
+    [removeDownload]
+  )
+
+  const onFileDownloadSuccess = useCallback(
+    async (state: DownloadState) => {
+      console.debug('onFileDownloadSuccess', state)
+      removeDownload(state.modelId)
+    },
+    [removeDownload]
+  )
+
+  useEffect(() => {
+    console.debug('DownloadListener: registering event listeners...')
+    events.on(DownloadEvent.onFileDownloadUpdate, onFileDownloadUpdate)
+    events.on(DownloadEvent.onFileDownloadError, onFileDownloadError)
+    events.on(DownloadEvent.onFileDownloadSuccess, onFileDownloadSuccess)
+    events.on(DownloadEvent.onFileDownloadStopped, onFileDownloadStopped)
+
+    return () => {
+      console.debug('DownloadListener: unregistering event listeners...')
+      events.off(DownloadEvent.onFileDownloadUpdate, onFileDownloadUpdate)
+      events.off(DownloadEvent.onFileDownloadError, onFileDownloadError)
+      events.off(DownloadEvent.onFileDownloadSuccess, onFileDownloadSuccess)
+      events.off(DownloadEvent.onFileDownloadStopped, onFileDownloadStopped)
+    }
+  }, [
+    onFileDownloadUpdate,
+    onFileDownloadError,
+    onFileDownloadSuccess,
+    onFileDownloadStopped,
+  ])
+
+  function renderGB(bytes: number): string {
+    const gb = bytes / 1024 ** 3
+    return ((gb * 100) / 100).toFixed(2)
+  }
+
  return (
    <Popover>
-      <PopoverTrigger>
-        <div className="bg-left-panel-fg/10 hover:bg-left-panel-fg/12 p-2 rounded-md my-1 relative border border-left-panel-fg/10 cursor-pointer text-left">
-          <div className="bg-primary font-bold size-5 rounded-full absolute -top-2 -right-1 flex items-center justify-center text-primary-fg">
-            2
+      {downloadCount > 0 && (
+        <PopoverTrigger>
+          <div className="bg-left-panel-fg/10 hover:bg-left-panel-fg/12 p-2 rounded-md my-1 relative border border-left-panel-fg/10 cursor-pointer text-left">
+            <div className="bg-primary font-bold size-5 rounded-full absolute -top-2 -right-1 flex items-center justify-center text-primary-fg">
+              {downloadCount}
+            </div>
+            <p className="text-left-panel-fg/80 font-medium">Downloads</p>
+            <div className="mt-2 flex items-center justify-between space-x-2">
+              <Progress value={overallProgress * 100} />
+              <span className="text-xs font-medium text-main-view-fg/80 shrink-0">
+                {overallProgress.toFixed(2)}%
+              </span>
+            </div>
          </div>
-          <p className="text-left-panel-fg/80 font-medium">Downloads</p>
-          <div className="mt-2 flex items-center justify-between space-x-2">
-            <Progress value={20} />
-            <span className="text-xs font-medium text-main-view-fg/80 shrink-0">
-              20%
-            </span>
-          </div>
-        </div>
-      </PopoverTrigger>
+        </PopoverTrigger>
+      )}
      <PopoverContent
        side="right"
        align="end"
@ -34,51 +131,33 @@ export function DownloadManagement() {
            <p className="text-xs text-main-view-fg/70">Downloading</p>
          </div>
          <div className="p-2 max-h-[300px] overflow-y-auto space-y-2">
-            <div className="bg-main-view-fg/4 rounded-md p-2">
-              <div className="flex items-center justify-between">
-                <p className="truncate text-main-view-fg/80">llama3.2:1b</p>
-                <div className="shrink-0 flex items-center space-x-0.5">
-                  <IconPlayerPauseFilled
-                    size={16}
-                    className="text-main-view-fg/70 cursor-pointer"
-                    title="Pause download"
-                  />
-                  <IconX
-                    size={16}
-                    className="text-main-view-fg/70 cursor-pointer"
-                    title="Cancel download"
-                  />
+            {downloadProcesses.map((download) => (
+              <div className="bg-main-view-fg/4 rounded-md p-2">
+                <div className="flex items-center justify-between">
+                  <p className="truncate text-main-view-fg/80">
+                    {download.name}
+                  </p>
+                  <div className="shrink-0 flex items-center space-x-0.5">
+                    {/* <IconPlayerPauseFilled
+                      size={16}
+                      className="text-main-view-fg/70 cursor-pointer"
+                      title="Pause download"
+                    /> */}
+                    <IconX
+                      size={16}
+                      className="text-main-view-fg/70 cursor-pointer"
+                      title="Cancel download"
+                      onClick={() => abortDownload(download.name)}
+                    />
+                  </div>
                </div>
-              </div>
-              <Progress value={25} className="my-2" />
-              <p className="text-main-view-fg/60 text-xs">
-                1065.28 MB/4.13 GB (25%)
-              </p>
-            </div>
-
-            <div className="bg-main-view-fg/4 rounded-md p-2">
-              <div className="flex items-center justify-between">
-                <p className="truncate text-main-view-fg/80">
-                  deepseek-r1:1.5b
+                <Progress value={download.progress * 100} className="my-2" />
+                <p className="text-main-view-fg/60 text-xs">
+                  {`${renderGB(download.current)} / ${renderGB(download.total)}`}{' '}
+                  GB ({download.progress.toFixed(2)}%)
                </p>
-                <div className="shrink-0 flex items-center space-x-0.5">
-                  <IconPlayerPauseFilled
-                    size={16}
-                    className="text-main-view-fg/70 cursor-pointer"
-                    title="Pause download"
-                  />
-                  <IconX
-                    size={16}
-                    className="text-main-view-fg/70 cursor-pointer"
-                    title="Cancel download"
-                  />
-                </div>
              </div>
-              <Progress value={80} className="my-2" />
-              <p className="text-main-view-fg/60 text-xs">
-                1065.28 MB/4.13 GB (80%)
-              </p>
-            </div>
+            ))}
          </div>
        </div>
      </PopoverContent>
--- a/web-app/src/hooks/useDownloadStore.ts
+++ b/web-app/src/hooks/useDownloadStore.ts
@ -0,0 +1,49 @@
+import { create } from 'zustand'
+
+export interface DownloadProgressProps {
+  id: string
+  progress: number
+  name: string
+  current: number
+  total: number
+}
+
+// Zustand store for thinking block state
+export type DownloadState = {
+  downloads: { [id: string]: DownloadProgressProps }
+  removeDownload: (id: string) => void
+  updateProgress: (
+    id: string,
+    progress: number,
+    name?: string,
+    current?: number,
+    total?: number
+  ) => void
+}
+
+/**
+ * This store is used to manage the download progress of files.
+ */
+export const useDownloadStore = create<DownloadState>((set) => ({
+  downloads: {},
+  removeDownload: (id: string) =>
+    set((state) => {
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      const { [id]: _, ...rest } = state.downloads
+      return { downloads: rest }
+    }),
+
+  updateProgress: (id, progress, name, current, total) =>
+    set((state) => ({
+      downloads: {
+        ...state.downloads,
+        [id]: {
+          ...state.downloads[id],
+          name: name || state.downloads[id]?.name || '',
+          progress,
+          current: current || state.downloads[id]?.current || 0,
+          total: total || state.downloads[id]?.total || 0,
+        },
+      },
+    })),
+}))
--- a/web-app/src/lib/completion.ts
+++ b/web-app/src/lib/completion.ts
@ -20,6 +20,7 @@ import { normalizeProvider } from './models'
 import { MCPTool } from '@/types/completion'
 import { CompletionMessagesBuilder } from './messages'
 import { ChatCompletionMessageToolCall } from 'openai/resources'
+import { callTool } from '@/services/mcp'

 /**
 * @fileoverview Helper functions for creating thread content.
@ -224,8 +225,7 @@ export const extractToolCall = (
      }

      if (deltaToolCalls[0]?.function?.arguments) {
-        currentCall!.function.arguments +=
-          deltaToolCalls[0].function.arguments
+        currentCall!.function.arguments += deltaToolCalls[0].function.arguments
      }
    }
  }
@ -268,7 +268,7 @@ export const postMessageProcessing = async (
        ],
      }

-      const result = await window.core.api.callTool({
+      const result = await callTool({
        toolName: toolCall.function.name,
        arguments: toolCall.function.arguments.length
          ? JSON.parse(toolCall.function.arguments)
--- a/web-app/src/lib/models.ts
+++ b/web-app/src/lib/models.ts
@ -65,5 +65,5 @@ export const extractModelRepo = (model?: string) => {
 */
 export const normalizeProvider = (provider: string) => {
  // TODO: After migrating to the new provider extension, remove this function
-  return provider === 'llama.cpp' ? 'llama-cpp' : provider
+  return provider === 'llama.cpp' ? 'cortex' : provider
 }
--- a/web-app/src/providers/ExtensionProvider.tsx
+++ b/web-app/src/providers/ExtensionProvider.tsx
@ -1,5 +1,6 @@
 import { ExtensionManager } from '@/lib/extension'
 import { APIs } from '@/lib/service'
+import { EventEmitter } from '@/services/eventsService'
 import { EngineManager, ModelManager } from '@janhq/core'
 import { PropsWithChildren, useCallback, useEffect, useState } from 'react'

@ -10,6 +11,7 @@ export function ExtensionProvider({ children }: PropsWithChildren) {
      api: APIs,
    }

+    window.core.events = new EventEmitter()
    window.core.extensionManager = new ExtensionManager()
    window.core.engineManager = new EngineManager()
    window.core.modelManager = new ModelManager()
--- a/web-app/src/routes/hub.tsx
+++ b/web-app/src/routes/hub.tsx
@ -15,6 +15,7 @@ import {
  DropdownMenuItem,
  DropdownMenuTrigger,
 } from '@/components/ui/dropdown-menu'
+import { downloadModel } from '@/services/models'

 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export const Route = createFileRoute(route.hub as any)({
@ -153,7 +154,13 @@ function Hub() {
                              <span className="text-main-view-fg/70 font-medium text-xs">
                                {toGigabytes(model.models?.[0]?.size)}
                              </span>
-                              <Button>Download</Button>
+                              <Button
+                                onClick={() =>
+                                  downloadModel(model.models[0]?.id)
+                                }
+                              >
+                                Download
+                              </Button>
                            </div>
                          </div>
                        }
@ -232,6 +239,9 @@ function Hub() {
                                        <div
                                          className="size-6 cursor-pointer flex items-center justify-center rounded hover:bg-main-view-fg/10 transition-all duration-200 ease-in-out"
                                          title="Edit All Servers JSON"
+                                          onClick={() =>
+                                            downloadModel(variant.id)
+                                          }
                                        >
                                          <IconDownload
                                            size={16}
--- a/web-app/src/services/eventsService.ts
+++ b/web-app/src/services/eventsService.ts
@ -0,0 +1,42 @@
+/* eslint-disable @typescript-eslint/no-unsafe-function-type */
+export class EventEmitter {
+  private handlers: Map<string, Function[]>
+
+  constructor() {
+    this.handlers = new Map<string, Function[]>()
+  }
+
+  public on(eventName: string, handler: Function): void {
+    if (!this.handlers.has(eventName)) {
+      this.handlers.set(eventName, [])
+    }
+
+    this.handlers.get(eventName)?.push(handler)
+  }
+
+  public off(eventName: string, handler: Function): void {
+    if (!this.handlers.has(eventName)) {
+      return
+    }
+
+    const handlers = this.handlers.get(eventName)
+    const index = handlers?.indexOf(handler)
+
+    if (index !== undefined && index !== -1) {
+      handlers?.splice(index, 1)
+    }
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  public emit(eventName: string, args: any): void {
+    if (!this.handlers.has(eventName)) {
+      return
+    }
+
+    const handlers = this.handlers.get(eventName)
+
+    handlers?.forEach((handler) => {
+      handler(args)
+    })
+  }
+}
--- a/web-app/src/services/mcp.ts
+++ b/web-app/src/services/mcp.ts
@ -1,10 +1,32 @@
+import { MCPTool } from '@/types/completion'

 /**
 * @description This file contains the functions to interact with the MCP API.
 * It includes functions to get and update the MCP configuration.
- * @param configs 
+ * @param configs
 */
 export const updateMCPConfig = async (configs: string) => {
  await window.core?.api?.saveMcpConfigs({ configs })
  await window.core?.api?.restartMcpServers()
 }
+
+/**
+ * @description This function gets the MCP configuration.
+ * @returns {Promise<string>} The MCP configuration.
+ */
+export const getTools = (): Promise<MCPTool[]> => {
+  return window.core?.api?.getTools()
+}
+
+/**
+ * @description This function invoke an MCP tool
+ * @param tool
+ * @param params
+ * @returns
+ */
+export const callTool = (args: {
+  toolName: string
+  arguments: object
+}): Promise<unknown> => {
+  return window.core?.api?.callTool(args)
+}
--- a/web-app/src/services/models.ts
+++ b/web-app/src/services/models.ts
@ -97,3 +97,43 @@ export const updateModel = async (
    throw error
  }
 }
+
+/**
+ * Downloads a model.
+ * @param model The model to download.
+ * @returns A promise that resolves when the model download task is created.
+ */
+export const downloadModel = async (id: string) => {
+  const extension = ExtensionManager.getInstance().get<ModelExtension>(
+    ExtensionTypeEnum.Model
+  )
+
+  if (!extension) throw new Error('Model extension not found')
+
+  try {
+    return await extension.pullModel(id)
+  } catch (error) {
+    console.error('Failed to download model:', error)
+    throw error
+  }
+}
+
+/**
+ * Aborts a model download.
+ * @param id
+ * @returns
+ */
+export const abortDownload = async (id: string) => {
+  const extension = ExtensionManager.getInstance().get<ModelExtension>(
+    ExtensionTypeEnum.Model
+  )
+
+  if (!extension) throw new Error('Model extension not found')
+
+  try {
+    return await extension.cancelModelPull(id)
+  } catch (error) {
+    console.error('Failed to abort model download:', error)
+    throw error
+  }
+}