feat: Add generation info to messages

This commit introduces a new feature that displays generation information for each message in the chat. The generation info is displayed in a popover and includes details about the model used, the prompt, and other relevant information. This helps users understand how their messages were generated and troubleshoot any issues that may arise. The generation info is retrieved from the LLM response and is stored in the database alongside other message details. This commit also includes translations for the generation info label in all supported languages.
2024-11-09 15:17:59 +05:30 · 2024-11-09 15:17:59 +05:30 · 9f383a81b6
commit 9f383a81b6
parent 9ecc8c4e75
26 changed files with 283 additions and 64 deletions
--- a/src/assets/locale/da/common.json
+++ b/src/assets/locale/da/common.json
@ -112,5 +112,6 @@
        "older": "Ældre"
    },
    "pin": "Fastgør",
-    "unpin": "Frigør"
+    "unpin": "Frigør",
+    "generationInfo": "Genererings Info"
 }
--- a/src/assets/locale/de/common.json
+++ b/src/assets/locale/de/common.json
@ -112,5 +112,6 @@
        "older": "Älter"
    },
    "pin": "Anheften",
-    "unpin": "Losheften"
+    "unpin": "Losheften",
+    "generationInfo": "Generierungsinformationen"
 }
--- a/src/assets/locale/en/common.json
+++ b/src/assets/locale/en/common.json
@ -116,5 +116,6 @@
        "older": "Older"
    },
    "pin": "Pin",
-    "unpin": "Unpin"
+    "unpin": "Unpin",
+    "generationInfo": "Generation Info"
 }
--- a/src/assets/locale/es/common.json
+++ b/src/assets/locale/es/common.json
@ -111,5 +111,6 @@
        "older": "Más antiguo"
    },
    "pin": "Fijar",
-    "unpin": "Desfijar"
+    "unpin": "Desfijar",
+    "generationInfo": "Información de Generación"
 }
--- a/src/assets/locale/fa/common.json
+++ b/src/assets/locale/fa/common.json
@ -105,5 +105,6 @@
        "older": "قدیمی‌تر"
    },
    "pin": "پین کردن",
-    "unpin": "حذف پین"
+    "unpin": "حذف پین",
+    "generationInfo": "اطلاعات تولید"
 }
--- a/src/assets/locale/fr/common.json
+++ b/src/assets/locale/fr/common.json
@ -111,5 +111,6 @@
        "older": "Plus ancien"
    },
    "pin": "Épingler",
-    "unpin": "Désépingler"
+    "unpin": "Désépingler",
+    "generationInfo": "Informations de génération"
 }
--- a/src/assets/locale/it/common.json
+++ b/src/assets/locale/it/common.json
@ -111,5 +111,6 @@
        "older": "Più Vecchi"
    },
    "pin": "Fissa",
-    "unpin": "Rimuovi"
+    "unpin": "Rimuovi",
+    "generationInfo": "Informazioni sulla Generazione"
 }
--- a/src/assets/locale/ja-JP/common.json
+++ b/src/assets/locale/ja-JP/common.json
@ -111,5 +111,6 @@
    "older": "それ以前"
  },
  "pin": "固定",
-  "unpin": "固定解除"
+  "unpin": "固定解除",
+  "generationInfo": "生成情報"
 }
--- a/src/assets/locale/ko/common.json
+++ b/src/assets/locale/ko/common.json
@ -111,5 +111,6 @@
    "older": "그 이전"
  },
  "pin": "고정",
-  "unpin": "고정 해제"
+  "unpin": "고정 해제",
+  "generationInfo": "생성 정보"
 }
--- a/src/assets/locale/ml/common.json
+++ b/src/assets/locale/ml/common.json
@ -110,5 +110,7 @@
        "older": "പഴയത്"
    },
    "pin": "പിൻ ചെയ്യുക",
-    "unpin": "അൺപിൻ ചെയ്യുക"
+    "unpin": "അൺപിൻ ചെയ്യുക",
+    "generationInfo": "ജനറേഷൻ വിവരങ്ങൾ"
+
 }
--- a/src/assets/locale/no/common.json
+++ b/src/assets/locale/no/common.json
@ -112,5 +112,6 @@
        "older": "Eldre"
    },
    "pin": "Fest",
-    "unpin": "Løsne"
+    "unpin": "Løsne",
+    "generationInfo": "Generasjonsinformasjon"
 }
--- a/src/assets/locale/pt-BR/common.json
+++ b/src/assets/locale/pt-BR/common.json
@ -111,5 +111,6 @@
    "older": "Mais Antigos"
  },
  "pin": "Fixar",
-  "unpin": "Desafixar"
+  "unpin": "Desafixar",
+  "generationInfo": "Informações de Geração"
 }
--- a/src/assets/locale/ru/common.json
+++ b/src/assets/locale/ru/common.json
@ -111,5 +111,6 @@
        "older": "Ранее"
    },
    "pin": "Закрепить",
-    "unpin": "Открепить"
+    "unpin": "Открепить",
+    "generationInfo": "Информация о генерации"
 }
--- a/src/assets/locale/sv/common.json
+++ b/src/assets/locale/sv/common.json
@ -116,5 +116,6 @@
        "older": "Äldre"
    },
    "pin": "Fäst",
-    "unpin": "Ta bort fäst"
+    "unpin": "Ta bort fäst",
+    "generationInfo": "Generationsinformation"
 }
--- a/src/assets/locale/zh/common.json
+++ b/src/assets/locale/zh/common.json
@ -111,5 +111,6 @@
        "older": "更早"
    },
    "pin": "置顶",
-    "unpin": "取消置顶"
+    "unpin": "取消置顶",
+    "generationInfo": "生成信息"
 }
--- a/src/components/Common/Playground/GenerationInfo.tsx
+++ b/src/components/Common/Playground/GenerationInfo.tsx
@ -0,0 +1,65 @@
+type GenerationMetrics = {
+  total_duration?: number
+  load_duration?: number
+  prompt_eval_count?: number
+  prompt_eval_duration?: number
+  eval_count?: number
+  eval_duration?: number
+  context?: string
+  response?: string
+}
+
+type Props = {
+  generationInfo: GenerationMetrics
+}
+
+export const GenerationInfo = ({ generationInfo }: Props) => {
+  if (!generationInfo) return null
+
+  const calculateTokensPerSecond = (
+    evalCount?: number,
+    evalDuration?: number
+  ) => {
+    if (!evalCount || !evalDuration) return 0
+    return (evalCount / evalDuration) * 1e9
+  }
+
+  const formatDuration = (nanoseconds?: number) => {
+    if (!nanoseconds) return "0ms"
+    const ms = nanoseconds / 1e6
+    if (ms < 1) return `${ms.toFixed(3)}ms`
+    if (ms < 1000) return `${Math.round(ms)}ms`
+    return `${(ms / 1000).toFixed(2)}s`
+  }
+
+  const metricsToDisplay = {
+    ...generationInfo,
+    ...(generationInfo?.eval_count && generationInfo?.eval_duration
+      ? {
+          tokens_per_second: calculateTokensPerSecond(
+            generationInfo.eval_count,
+            generationInfo.eval_duration
+          ).toFixed(2)
+        }
+      : {})
+  }
+
+  return (
+    <div className="p-2 w-full">
+      <div className="flex flex-col gap-2">
+        {Object.entries(metricsToDisplay)
+          .filter(([key]) => key !== "model")
+          .map(([key, value]) => (
+            <div key={key} className="flex flex-wrap justify-between">
+              <div className="font-medium text-xs">{key}</div>
+              <div className="font-medium text-xs break-all">
+                {key.includes("duration")
+                  ? formatDuration(value as number)
+                  : String(value)}
+              </div>
+            </div>
+          ))}
+      </div>
+    </div>
+  )
+}
--- a/src/components/Common/Playground/Message.tsx
+++ b/src/components/Common/Playground/Message.tsx
@ -1,10 +1,11 @@
 import Markdown from "../../Common/Markdown"
 import React from "react"
-import { Tag, Image, Tooltip, Collapse } from "antd"
+import { Tag, Image, Tooltip, Collapse, Popover } from "antd"
 import { WebSearch } from "./WebSearch"
 import {
  CheckIcon,
  ClipboardIcon,
+  InfoIcon,
  Pen,
  PlayIcon,
  RotateCcw,
@ -16,6 +17,7 @@ import { MessageSource } from "./MessageSource"
 import { useTTS } from "@/hooks/useTTS"
 import { tagColors } from "@/utils/color"
 import { removeModelSuffix } from "@/db/models"
+import { GenerationInfo } from "./GenerationInfo"

 type Props = {
  message: string
@ -37,6 +39,7 @@ type Props = {
  hideEditAndRegenerate?: boolean
  onSourceClick?: (source: any) => void
  isTTSEnabled?: boolean
+  generationInfo?: any
 }

 export const PlaygroundMessage = (props: Props) => {
@ -206,6 +209,18 @@ export const PlaygroundMessage = (props: Props) => {
                      </Tooltip>
                    )}

+                    {props.generationInfo && (
+                      <Popover
+                        content={
+                         <GenerationInfo generationInfo={props.generationInfo} />
+                        }
+                        title={t("generationInfo")}>
+                        <button className="flex items-center justify-center w-6 h-6 rounded-full bg-gray-100 dark:bg-gray-800 hover:bg-gray-200 dark:hover:bg-gray-700 transition-colors duration-200 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-gray-500">
+                          <InfoIcon className="w-3 h-3 text-gray-400 group-hover:text-gray-500" />
+                        </button>
+                      </Popover>
+                    )}
+
                    {!props.hideEditAndRegenerate &&
                      props.currentMessageIndex === props.totalMessages - 1 && (
                        <Tooltip title={t("regenerate")}>
--- a/src/components/Option/Playground/PlaygroundChat.tsx
+++ b/src/components/Option/Playground/PlaygroundChat.tsx
@ -54,6 +54,7 @@ export const PlaygroundChat = () => {
              setIsSourceOpen(true)
            }}
            isTTSEnabled={ttsEnabled}
+            generationInfo={message?.generationInfo}
          />
        ))}
        {messages.length > 0 && (
--- a/src/components/Sidepanel/Chat/body.tsx
+++ b/src/components/Sidepanel/Chat/body.tsx
@ -47,6 +47,7 @@ export const SidePanelBody = () => {
            setIsSourceOpen(true)
          }}
          isTTSEnabled={ttsEnabled}
+          generationInfo={message?.generationInfo}
        />
      ))}
      <div className="w-full h-48 flex-shrink-0"></div>
--- a/src/db/index.ts
+++ b/src/db/index.ts
@ -33,6 +33,7 @@ type Message = {
  search?: WebSearch
  createdAt: number
  messageType?: string
+  generationInfo?: any
 }

 type Webshare = {
@ -254,7 +255,8 @@ export const saveMessage = async (
  images: string[],
  source?: any[],
  time?: number,
-  message_type?: string
+  message_type?: string,
+  generationInfo?: any
 ) => {
  const id = generateID()
  let createdAt = Date.now()
@ -270,7 +272,8 @@ export const saveMessage = async (
    images,
    createdAt,
    sources: source,
-    messageType: message_type
+    messageType: message_type,
+    generationInfo: generationInfo
  }
  const db = new PageAssitDatabase()
  await db.addMessage(message)
--- a/src/hooks/chat-helper/index.ts
+++ b/src/hooks/chat-helper/index.ts
@ -118,7 +118,7 @@ export const saveMessageOnSuccess = async ({
  fullText,
  source,
  message_source = "web-ui",
-  message_type
+  message_type, generationInfo
 }: {
  historyId: string | null
  setHistoryId: (historyId: string) => void
@ -130,6 +130,7 @@ export const saveMessageOnSuccess = async ({
  source: any[]
  message_source?: "copilot" | "web-ui",
  message_type?: string
+  generationInfo?: any
 }) => {
  if (historyId) {
    if (!isRegenerate) {
@ -141,7 +142,8 @@ export const saveMessageOnSuccess = async ({
        [image],
        [],
        1,
-        message_type
+        message_type,
+        generationInfo
      )
    }
    await saveMessage(
@ -152,7 +154,8 @@ export const saveMessageOnSuccess = async ({
      [],
      source,
      2,
-      message_type
+      message_type,
+      generationInfo
    )
    await setLastUsedChatModel(historyId, selectedModel!)
  } else {
@ -166,7 +169,8 @@ export const saveMessageOnSuccess = async ({
      [image],
      [],
      1,
-      message_type
+      message_type,
+      generationInfo
    )
    await saveMessage(
      newHistoryId.id,
@ -176,7 +180,8 @@ export const saveMessageOnSuccess = async ({
      [],
      source,
      2,
-      message_type
+      message_type,
+      generationInfo
    )
    setHistoryId(newHistoryId.id)
    await setLastUsedChatModel(newHistoryId.id, selectedModel!)
--- a/src/hooks/useMessage.tsx
+++ b/src/hooks/useMessage.tsx
@ -328,10 +328,25 @@ export const useMessage = () => {

      const applicationChatHistory = generateHistory(history, selectedModel)

+      let generationInfo: any | undefined = undefined
+
      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
-          signal: signal
+          signal: signal,
+          callbacks: [
+            {
+              handleLLMEnd(
+                output: any,
+              ): any {
+                try {
+                  generationInfo = output?.generations?.[0][0]?.generationInfo
+                } catch (e) {
+                  console.log("handleLLMEnd error", e)
+                }
+              }
+            }
+          ]
        }
      )
      let count = 0
@ -361,7 +376,8 @@ export const useMessage = () => {
            return {
              ...message,
              message: fullText,
-              sources: source
+              sources: source,
+              generationInfo
            }
          }
          return message
@ -390,7 +406,8 @@ export const useMessage = () => {
        image,
        fullText,
        source,
-        message_source: "copilot"
+        message_source: "copilot",
+        generationInfo
      })

      setIsProcessing(false)
@ -544,10 +561,25 @@ export const useMessage = () => {
        )
      }

+      let generationInfo: any | undefined = undefined
+
      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
-          signal: signal
+          signal: signal,
+          callbacks: [
+            {
+              handleLLMEnd(
+                output: any,
+              ): any {
+                try {
+                  generationInfo = output?.generations?.[0][0]?.generationInfo
+                } catch (e) {
+                  console.log("handleLLMEnd error", e)
+                }
+              }
+            }
+          ]
        }
      )
      let count = 0
@ -576,7 +608,8 @@ export const useMessage = () => {
          if (message.id === generateMessageId) {
            return {
              ...message,
-              message: fullText
+              message: fullText,
+              generationInfo
            }
          }
          return message
@ -605,7 +638,8 @@ export const useMessage = () => {
        image,
        fullText,
        source: [],
-        message_source: "copilot"
+        message_source: "copilot",
+        generationInfo
      })

      setIsProcessing(false)
@ -789,10 +823,24 @@ export const useMessage = () => {
        )
      }

+      let generationInfo: any | undefined = undefined
      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
-          signal: signal
+          signal: signal,
+          callbacks: [
+            {
+              handleLLMEnd(
+                output: any,
+              ): any {
+                try {
+                  generationInfo = output?.generations?.[0][0]?.generationInfo
+                } catch (e) {
+                  console.log("handleLLMEnd error", e)
+                }
+              }
+            }
+          ]
        }
      )
      let count = 0
@ -822,7 +870,8 @@ export const useMessage = () => {
            return {
              ...message,
              message: fullText,
-              sources: source
+              sources: source,
+              generationInfo
            }
          }
          return message
@ -850,7 +899,8 @@ export const useMessage = () => {
        message,
        image,
        fullText,
-        source
+        source,
+        generationInfo
      })

      setIsProcessing(false)
@ -982,8 +1032,23 @@ export const useMessage = () => {
        })
      }

+      let generationInfo: any | undefined = undefined
+
      const chunks = await ollama.stream([humanMessage], {
-        signal: signal
+        signal: signal,
+        callbacks: [
+          {
+            handleLLMEnd(
+              output: any,
+            ): any {
+              try {
+                generationInfo = output?.generations?.[0][0]?.generationInfo
+              } catch (e) {
+                console.log("handleLLMEnd error", e)
+              }
+            }
+          }
+        ]
      })
      let count = 0
      for await (const chunk of chunks) {
@ -1011,7 +1076,8 @@ export const useMessage = () => {
          if (message.id === generateMessageId) {
            return {
              ...message,
-              message: fullText
+              message: fullText,
+              generationInfo
            }
          }
          return message
@ -1042,7 +1108,8 @@ export const useMessage = () => {
        fullText,
        source: [],
        message_source: "copilot",
-        message_type: messageType
+        message_type: messageType,
+        generationInfo
      })

      setIsProcessing(false)
--- a/src/hooks/useMessageOption.tsx
+++ b/src/hooks/useMessageOption.tsx
@ -243,10 +243,23 @@ export const useMessageOption = () => {
        )
      }

+      let generationInfo: any | undefined = undefined
+
      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
-          signal: signal
+          signal: signal,
+          callbacks: [
+            {
+              handleLLMEnd(output: any): any {
+                try {
+                  generationInfo = output?.generations?.[0][0]?.generationInfo
+                } catch (e) {
+                  console.log("handleLLMEnd error", e)
+                }
+              }
+            }
+          ]
        }
      )
      let count = 0
@ -276,7 +289,8 @@ export const useMessageOption = () => {
            return {
              ...message,
              message: fullText,
-              sources: source
+              sources: source,
+              generationInfo
            }
          }
          return message
@ -304,7 +318,8 @@ export const useMessageOption = () => {
        message,
        image,
        fullText,
-        source
+        source,
+        generationInfo
      })

      setIsProcessing(false)
@ -465,10 +480,23 @@ export const useMessageOption = () => {
        )
      }

+      let generationInfo: any | undefined = undefined
+
      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
-          signal: signal
+          signal: signal,
+          callbacks: [
+            {
+              handleLLMEnd(output: any): any {
+                try {
+                  generationInfo = output?.generations?.[0][0]?.generationInfo
+                } catch (e) {
+                  console.log("handleLLMEnd error", e)
+                }
+              }
+            }
+          ],
        }
      )

@ -498,7 +526,8 @@ export const useMessageOption = () => {
          if (message.id === generateMessageId) {
            return {
              ...message,
-              message: fullText
+              message: fullText,
+              generationInfo
            }
          }
          return message
@ -526,7 +555,8 @@ export const useMessageOption = () => {
        message,
        image,
        fullText,
-        source: []
+        source: [],
+        generationInfo
      })

      setIsProcessing(false)
@ -711,10 +741,23 @@ export const useMessageOption = () => {

      const applicationChatHistory = generateHistory(history, selectedModel)

+      let generationInfo: any | undefined = undefined
+
      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
-          signal: signal
+          signal: signal,
+          callbacks: [
+            {
+              handleLLMEnd(output: any): any {
+                try {
+                  generationInfo = output?.generations?.[0][0]?.generationInfo
+                } catch (e) {
+                  console.log("handleLLMEnd error", e)
+                }
+              }
+            }
+          ]
        }
      )
      let count = 0
@ -744,7 +787,8 @@ export const useMessageOption = () => {
            return {
              ...message,
              message: fullText,
-              sources: source
+              sources: source,
+              generationInfo
            }
          }
          return message
@ -772,7 +816,8 @@ export const useMessageOption = () => {
        message,
        image,
        fullText,
-        source
+        source,
+        generationInfo
      })

      setIsProcessing(false)
--- a/src/models/index.ts
+++ b/src/models/index.ts
@ -49,7 +49,7 @@ export const pageAssistModel = async ({
      configuration: {
        apiKey: providerInfo.apiKey || "temp",
        baseURL: providerInfo.baseUrl || "",
-      }
+      },
    }) as any
  }

--- a/src/types/message.ts
+++ b/src/types/message.ts
@ -1,19 +1,20 @@
 type WebSearch = {
-    search_engine: string
-    search_url: string
-    search_query: string
-    search_results: {
-      title: string
-      link: string
-    }[]
-  }
-  export type Message = {
-    isBot: boolean
-    name: string
-    message: string
-    sources: any[]
-    images?: string[]
-    search?: WebSearch
-    messageType?: string
-    id?: string
-  }
+  search_engine: string
+  search_url: string
+  search_query: string
+  search_results: {
+    title: string
+    link: string
+  }[]
+}
+export type Message = {
+  isBot: boolean
+  name: string
+  message: string
+  sources: any[]
+  images?: string[]
+  search?: WebSearch
+  messageType?: string
+  id?: string
+  generationInfo?: any
+}
--- a/wxt.config.ts
+++ b/wxt.config.ts
@ -50,7 +50,7 @@ export default defineConfig({
  outDir: "build",

  manifest: {
-    version: "1.3.3",
+    version: "1.3.4",
    name:
      process.env.TARGET === "firefox"
        ? "Page Assist - A Web UI for Local AI Models"