chore: streaming tool output (#5237)

* enhancement: tool streaming output * chore: update memo * fix: streaming * chore: update stream tools arguments * chore: update condition * fix: style * fix: style * chore: fix stop button * chore: update color accent and hide arrow button --------- Co-authored-by: Louis <louis@jan.ai>
2025-06-11 14:35:41 +07:00 · 2025-06-11 14:35:41 +07:00 · 808fdb02a7
commit 808fdb02a7
parent cfefcb00cb
9 changed files with 144 additions and 52 deletions
--- a/web-app/src/containers/RenderMarkdown.tsx
+++ b/web-app/src/containers/RenderMarkdown.tsx
@ -22,6 +22,7 @@ interface MarkdownProps {
  components?: Components
  enableRawHtml?: boolean
  isUser?: boolean
+  isWrapping?: boolean
 }

 function RenderMarkdownComponent({
@ -30,6 +31,7 @@ function RenderMarkdownComponent({
  className,
  isUser,
  components,
+  isWrapping,
 }: MarkdownProps) {
  const { codeBlockStyle, showLineNumbers } = useCodeblock()

@ -117,9 +119,13 @@ function RenderMarkdownComponent({
              showLineNumbers={showLineNumbers}
              wrapLines={true}
              // Temporary comment we try calculate main area width on __root
-              // lineProps={{
-              //   style: { wordBreak: 'break-all', whiteSpace: 'pre-wrap' },
-              // }}
+              lineProps={
+                isWrapping
+                  ? {
+                      style: { wordBreak: 'break-all', whiteSpace: 'pre-wrap' },
+                    }
+                  : {}
+              }
              customStyle={{
                margin: 0,
                padding: '8px',
--- a/web-app/src/containers/StreamingContent.tsx
+++ b/web-app/src/containers/StreamingContent.tsx
@ -32,6 +32,12 @@ export const StreamingContent = memo(({ threadId }: Props) => {
    return extractReasoningSegment(text)
  }, [streamingContent])

+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const streamingTools: any = useMemo(() => {
+    const calls = streamingContent?.metadata?.tool_calls
+    return calls
+  }, [streamingContent])
+
  const lastAssistant = useMemo(() => {
    return [...messages].reverse().find((m) => m.role === 'assistant')
  }, [messages])
@ -52,6 +58,14 @@ export const StreamingContent = memo(({ threadId }: Props) => {
  // The streaming content is always the last message
  return (
    <ThreadContent
+      streamTools={{
+        tool_calls: {
+          function: {
+            name: streamingTools?.[0]?.function?.name as string,
+            arguments: streamingTools?.[0]?.function?.arguments as string,
+          },
+        },
+      }}
      {...streamingContent}
      isLastMessage={true}
      showAssistant={
--- a/web-app/src/containers/ThreadContent.tsx
+++ b/web-app/src/containers/ThreadContent.tsx
@ -77,6 +77,8 @@ export const ThreadContent = memo(
      isLastMessage?: boolean
      index?: number
      showAssistant?: boolean
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      streamTools?: any
    }
  ) => {
    const [message, setMessage] = useState(item.content?.[0]?.text?.value || '')
@ -316,8 +318,17 @@ export const ThreadContent = memo(
                {(item.metadata.tool_calls as ToolCall[]).map((toolCall) => (
                  <ToolCallBlock
                    id={toolCall.tool?.id ?? 0}
-                    name={toolCall.tool?.function?.name ?? ''}
                    key={toolCall.tool?.id}
+                    name={
+                      (item.streamTools?.tool_calls?.function?.name ||
+                        toolCall.tool?.function?.name) ??
+                      ''
+                    }
+                    args={
+                      item.streamTools?.tool_calls?.function?.arguments ||
+                      toolCall.tool?.function?.arguments ||
+                      undefined
+                    }
                    result={JSON.stringify(toolCall.response)}
                    loading={toolCall.state === 'pending'}
                  />
@ -332,7 +343,7 @@ export const ThreadContent = memo(
                    'flex items-center gap-2',
                    item.isLastMessage &&
                      streamingContent &&
-                      'opacity-0 visinility-hidden pointer-events-none'
+                      'opacity-0 visibility-hidden pointer-events-none'
                  )}
                >
                  <CopyButton text={item.content?.[0]?.text.value || ''} />
--- a/web-app/src/containers/ToolCallBlock.tsx
+++ b/web-app/src/containers/ToolCallBlock.tsx
@ -9,10 +9,12 @@ import {
  DialogHeader,
  DialogTitle,
 } from '@/components/ui/dialog'
+import { twMerge } from 'tailwind-merge'

 interface Props {
  result: string
  name: string
+  args: object
  id: number
  loading: boolean
 }
@ -90,7 +92,7 @@ const ContentItemRenderer = ({
  if (item.type === 'image' && item.data && item.mimeType) {
    const imageUrl = createDataUrl(item.data, item.mimeType)
    return (
-      <div key={index} className="mt-3">
+      <div key={index} className="my-3">
        <img
          src={imageUrl}
          alt={`Result image ${index + 1}`}
@ -123,9 +125,9 @@ const ContentItemRenderer = ({
  )
 }

-const ToolCallBlock = ({ id, name, result, loading }: Props) => {
+const ToolCallBlock = ({ id, name, result, loading, args }: Props) => {
  const { collapseState, setCollapseState } = useToolCallBlockStore()
-  const isExpanded = collapseState[id] ?? false
+  const isExpanded = collapseState[id] ?? (loading ? true : false)
  const [modalImage, setModalImage] = useState<{
    url: string
    alt: string
@ -145,10 +147,9 @@ const ToolCallBlock = ({ id, name, result, loading }: Props) => {
  }

  // Parse the MCP response and extract content items
-  const { parsedResult, contentItems, hasStructuredContent, parseError } =
-    useMemo(() => {
-      return parseMCPResponse(result)
-    }, [result])
+  const { parsedResult, contentItems, hasStructuredContent } = useMemo(() => {
+    return parseMCPResponse(result)
+  }, [result])

  return (
    <div
@ -158,17 +159,36 @@ const ToolCallBlock = ({ id, name, result, loading }: Props) => {
      <div className="rounded-lg bg-main-view-fg/4 border border-dashed border-main-view-fg/10">
        <div className="flex items-center gap-3 p-2" onClick={handleClick}>
          {loading && (
-            <Loader className="size-4 animate-spin text-main-view-fg/60" />
+            <div className="w-4 h-4">
+              <Loader className="size-4 animate-spin text-main-view-fg/60" />
+            </div>
          )}
          <button className="flex items-center gap-2 focus:outline-none">
-            {isExpanded ? (
-              <ChevronUp className="h-4 w-4" />
-            ) : (
-              <ChevronDown className="h-4 w-4" />
+            {!loading && (
+              <>
+                {isExpanded ? (
+                  <>
+                    <div className="ml-1 w-4 h-4">
+                      <ChevronUp className="h-4 w-4" />
+                    </div>
+                  </>
+                ) : (
+                  <div className="ml-1 w-4 h-4">
+                    <ChevronDown className="h-4 w-4" />
+                  </div>
+                )}
+              </>
            )}
            <span className="font-medium text-main-view-fg/80">
-              View result from{' '}
-              <span className="font-medium text-main-view-fg">{name}</span>
+              <span className="font-medium text-main-view-fg mr-2">{name}</span>
+              <span
+                className={twMerge(
+                  'text-xs bg-main-view-fg/4 rounded-sm p-1',
+                  loading ? 'text-main-view-fg/40' : 'text-accent'
+                )}
+              >
+                {loading ? 'Calling tool' : 'Completed'}{' '}
+              </span>
            </span>
          </button>
        </div>
@ -179,36 +199,43 @@ const ToolCallBlock = ({ id, name, result, loading }: Props) => {
            isExpanded ? '' : 'max-h-0 overflow-hidden'
          )}
        >
-          <div className="mt-2 text-main-view-fg/60">
-            {hasStructuredContent ? (
-              /* Render each content item individually based on its type */
-              <div className="space-y-2">
-                {contentItems.map((item, index) => (
-                  <ContentItemRenderer
-                    key={index}
-                    item={item}
-                    index={index}
-                    onImageClick={handleImageClick}
+          <div className="mt-2 text-main-view-fg/60 max-w-[89%] overflow-hidden">
+            {args && Object.keys(args).length > 3 && (
+              <>
+                <p className="mb-3">Arguments:</p>
+                <RenderMarkdown
+                  isWrapping={true}
+                  content={'```json\n' + args + '\n```'}
+                />
+              </>
+            )}
+
+            {result && (
+              <>
+                <p>Output:</p>
+                {hasStructuredContent ? (
+                  /* Render each content item individually based on its type */
+                  <div className="space-y-2">
+                    {contentItems.map((item, index) => (
+                      <ContentItemRenderer
+                        key={index}
+                        item={item}
+                        index={index}
+                        onImageClick={handleImageClick}
+                      />
+                    ))}
+                  </div>
+                ) : (
+                  /* Fallback: render as JSON for valid JSON but unstructured responses */
+                  <RenderMarkdown
+                    content={
+                      '```json\n' +
+                      JSON.stringify(parsedResult, null, 2) +
+                      '\n```'
+                    }
                  />
-                ))}
-              </div>
-            ) : parseError ? (
-              /* Handle JSON parse error - render as plain text */
-              <div className="mt-3 p-3 bg-main-view-fg/5 rounded-md border border-main-view-fg/10">
-                <div className="text-sm font-medium text-main-view-fg/80 mb-2">
-                  Raw Response:
-                </div>
-                <div className="whitespace-pre-wrap font-mono text-sm">
-                  {parsedResult as string}
-                </div>
-              </div>
-            ) : (
-              /* Fallback: render as JSON for valid JSON but unstructured responses */
-              <RenderMarkdown
-                content={
-                  '```json\n' + JSON.stringify(parsedResult, null, 2) + '\n```'
-                }
-              />
+                )}
+              </>
            )}
          </div>
        </div>
--- a/web-app/src/hooks/useAppState.ts
+++ b/web-app/src/hooks/useAppState.ts
@ -2,6 +2,7 @@ import { create } from 'zustand'
 import { ThreadMessage } from '@janhq/core'
 import { MCPTool } from '@/types/completion'
 import { useAssistant } from './useAssistant'
+import { ChatCompletionMessageToolCall } from 'openai/resources'

 type AppState = {
  streamingContent?: ThreadMessage
@ -10,8 +11,12 @@ type AppState = {
  serverStatus: 'running' | 'stopped' | 'pending'
  abortControllers: Record<string, AbortController>
  tokenSpeed?: TokenSpeed
+  currentToolCall?: ChatCompletionMessageToolCall
  setServerStatus: (value: 'running' | 'stopped' | 'pending') => void
  updateStreamingContent: (content: ThreadMessage | undefined) => void
+  updateCurrentToolCall: (
+    toolCall: ChatCompletionMessageToolCall | undefined
+  ) => void
  updateLoadingModel: (loading: boolean) => void
  updateTools: (tools: MCPTool[]) => void
  setAbortController: (threadId: string, controller: AbortController) => void
@ -26,6 +31,7 @@ export const useAppState = create<AppState>()((set) => ({
  serverStatus: 'stopped',
  abortControllers: {},
  tokenSpeed: undefined,
+  currentToolCall: undefined,
  updateStreamingContent: (content: ThreadMessage | undefined) => {
    set(() => ({
      streamingContent: content
@ -40,6 +46,11 @@ export const useAppState = create<AppState>()((set) => ({
        : undefined,
    }))
  },
+  updateCurrentToolCall: (toolCall) => {
+    set(() => ({
+      currentToolCall: toolCall,
+    }))
+  },
  updateLoadingModel: (loading) => {
    set({ loadingModel: loading })
  },
--- a/web-app/src/hooks/useChat.ts
+++ b/web-app/src/hooks/useChat.ts
@ -181,8 +181,21 @@ export const useChat = () => {
          } else {
            for await (const part of completion) {
              const delta = part.choices[0]?.delta?.content || ''
+
              if (part.choices[0]?.delta?.tool_calls) {
-                extractToolCall(part, currentCall, toolCalls)
+                const calls = extractToolCall(part, currentCall, toolCalls)
+                const currentContent = newAssistantThreadContent(
+                  activeThread.id,
+                  accumulatedText,
+                  {
+                    tool_calls: calls.map((e) => ({
+                      ...e,
+                      state: 'pending',
+                    })),
+                  }
+                )
+                updateStreamingContent(currentContent)
+                await new Promise((resolve) => setTimeout(resolve, 0))
              }
              if (delta) {
                accumulatedText += delta
@ -190,7 +203,13 @@ export const useChat = () => {
                // Use a timeout to prevent React from batching updates too quickly
                const currentContent = newAssistantThreadContent(
                  activeThread.id,
-                  accumulatedText
+                  accumulatedText,
+                  {
+                    tool_calls: toolCalls.map((e) => ({
+                      ...e,
+                      state: 'pending',
+                    })),
+                  }
                )
                updateStreamingContent(currentContent)
                updateTokenSpeed(currentContent)
@ -225,6 +244,7 @@ export const useChat = () => {
            allowAllMCPPermissions
          )
          addMessage(updatedMessage ?? finalContent)
+          updateStreamingContent(emptyThreadContent)
          updateThreadTimestamp(activeThread.id)

          isCompleted = !toolCalls.length
--- a/web-app/src/lib/completion.ts
+++ b/web-app/src/lib/completion.ts
@ -63,7 +63,8 @@ export const newUserThreadContent = (
 */
 export const newAssistantThreadContent = (
  threadId: string,
-  content: string
+  content: string,
+  metadata: Record<string, unknown> = {}
 ): ThreadMessage => ({
  type: 'text',
  role: ChatCompletionRole.Assistant,
@ -82,6 +83,7 @@ export const newAssistantThreadContent = (
  status: MessageStatus.Ready,
  created_at: 0,
  completed_at: 0,
+  metadata,
 })

 /**
--- a/web-app/src/routes/__root.tsx
+++ b/web-app/src/routes/__root.tsx
@ -43,7 +43,7 @@ const AppLayout = () => {
          <div
            className={cn(
              'h-full flex w-full p-1',
-              isLeftPanelOpen && 'w-[calc(100%-192px)]'
+              isLeftPanelOpen && 'w-[calc(100%-198px)]'
            )}
          >
            <div className="bg-main-view text-main-view-fg border border-main-view-fg/5 w-full rounded-lg overflow-hidden">
--- a/web-app/src/types/message.d.ts
+++ b/web-app/src/types/message.d.ts
@ -3,6 +3,7 @@ type ToolCall = {
    id?: number
    function?: {
      name?: string
+      arguments?: object
    }
  }
  response?: unknown