@stack-spot/ai-chat-widget 2.2.0 → 2.2.1-beta.1

package/src/views/Chat/StepsList.tsx

@@ -1,96 +1,389 @@
-import { Icon } from '@stack-spot/citric-icons'
-import { Button, ProgressCircular, Text } from '@stack-spot/citric-react'
+import { Accordion, Badge, Button, Card, Column, Divider, Icon, IconBox, ImageWithFallback, ProgressCircular, Row, Skeleton, Text } from '@stack-spot/citric-react'
 import { AnimatedHeight } from '@stack-spot/portal-components/AnimatedHeight'
-import { ChatStep, StepChatStep } from '@stack-spot/portal-network'
+import { ChatStep, StepChatStep, ToolChatStep } from '@stack-spot/portal-network'
 import { theme } from '@stack-spot/portal-theme'
 import { Dictionary, useTranslate } from '@stack-spot/portal-translate'
-import { findLastIndex } from 'lodash'
-import { useState } from 'react'
-import { useWidget } from '../../context/hooks'
+import { findLast, findLastIndex } from 'lodash'
+import React, { useEffect, useMemo } from 'react'
+import styled from 'styled-components'
+import { Markdown } from '../../components/Markdown'
+import { useChat, useChatMessages, useCurrentChat, useCurrentChatMessages, useWidget } from '../../context/hooks'
+import { ChatEntry } from '../../state/ChatEntry'
+import { planningToolDictionaryHelper } from '../../utils/planning-tool'
+import { updateToolStep } from '../../utils/update-tool-step'
+import { onCopyCode } from './events'
 
 interface Props {
   steps: ChatStep[],
   messageId: number,
   chatId: string,
+  userHasAlreadyAnswered?: boolean,
+}
+
+interface StepChatStepWithTarget extends Omit<StepChatStep, 'status' | 'id' | 'type'> {
+  status: 'pending' | 'running' | 'success' | 'error' | 'target' | 'awaiting_approval',
 }
 
 interface StepProps {
-  step: StepChatStep,
+  step: StepChatStepWithTarget,
   index: number,
-  total: number,
+  total?: number,
+  totalTools?: number,
+  isAllDone?: boolean,
   onClick?: () => void,
 }
 
-function getStatusIcon(status: ChatStep['status']) {
+function getStatusIcon(status: StepChatStepWithTarget['status'] | 'target', isDone?: boolean) {
   const iconProps = { style: { color: theme.color.light[700] }, size: 'xs' } as const
   switch (status) {
     case 'error': return <Icon group="fill" icon="TimesCircle" {...iconProps} />
     case 'success': return <Icon group="fill" icon="CheckCircle" {...iconProps} />
-    case 'pending': return <Icon icon="Spaces" {...iconProps} />
+    case 'pending': return <Icon group="fill" icon="Circle" {...iconProps} style={{ color: theme.color.light[600] }} />
+    case 'awaiting_approval': return <Icon group="fill" icon="ExclamationTriangle" {...iconProps} />
+    case 'target': return <Icon icon="Target" {...iconProps} style={{ color: isDone ? theme.color.light[700] : theme.color.light[600] }} />
     case 'running': return <ProgressCircular className="loading" colorScheme="inverse" size="xs" />
   }
 }
 
-const Step = ({ step, index, total, onClick }: StepProps) => {
+const StepAccordionHeader = ({ step, index, expand }: Pick<StepProps, 'step' | 'index'> & { expand: React.ReactElement }) => {
+  const t = useTranslate(dictionary)
+  return <Row gap="8px">
+    {expand}
+    {step.status === 'target' ? <Text className="step-title" appearance="body2" color="light.700">{t.planGoal}:</Text> :
+      <Badge colorScheme="inverse" appearance="square">
+        {t.step} {index}
+      </Badge>}
+    <Text className="step-title" appearance="body2">
+      {step.input}
+    </Text>
+    {step.status === 'awaiting_approval' &&
+      <Badge appearance="square" style={{ backgroundColor: theme.color.gray[800], color: theme.color.gray[50] }}>
+        <Icon icon="Security" />
+        {t.pendingReview}
+      </Badge>}
+  </Row>
+}
+
+const StyledCard = styled(Card)`
+  &:hover { 
+    background-color: ${theme.color.light[500]}
+  }
+`
+
+const Step = ({ step, index, onClick, total, totalTools, isAllDone }: StepProps) => {
   const t = useTranslate(dictionary)
+  const status = getStatusIcon(step.status, isAllDone)
+  const hasTools = step.attempts?.[0]?.tools && step.attempts?.[0]?.tools?.length > 0
+
   return (
     <li tabIndex={onClick ? 0 : undefined} onClick={onClick} role={onClick ? 'button' : 'listitem'}>
-      <div className="step-status-icon">{getStatusIcon(step.status)}</div>
-      <Text className="step-title" appearance="microtext1" color="light.700">
-        {t.step} {index}/{total}: {step.input}
-      </Text>
+      <Row gap="4px" alignItems="center">
+        <div className="step-status-icon">{status}</div>
+        <StyledCard p="8px" w="80%">
+          <Accordion header={expand => <StepAccordionHeader step={step} index={index} expand={expand} />}>
+            <Column pt="12px">
+              {total ?
+                <Row gap="40px">
+                  <Row gap="4px">
+                    <Icon icon="Hashtag" size="sm" color="light.700" />
+                    <Text color="light.700">{t.totalSteps}</Text>
+                    {total}
+                  </Row>
+                  <Row gap="4px">
+                    <Icon icon="BorderRadius" size="sm" color="light.700" />
+                    <Text color="light.700">{t.totalTools}</Text>
+                    {totalTools ?? 0}
+                  </Row>
+                </Row>
+                : <>
+                  <Row pb="8px">
+                    <Icon icon="Target" size="sm" color="light.700" />
+                    <Text color="light.700" tag="span" style={{ margin: '0 4px' }}>{t.stepGoal}:</Text>
+                    <Text tag="span">
+                      {step.input}
+                    </Text>
+                  </Row>
+                  {hasTools ? <>
+                    <Row gap="4px">
+                      <Icon icon="BorderRadius" size="sm" color="light.700" />
+                      <Text color="light.700">{t.toolsToBeExecuted}:</Text>
+                    </Row>
+                    <ul className="tools-list">
+                      {step.attempts?.[0]?.tools?.map((tool, index) => (<li key={`${tool.id}-${index}`}>
+                        <Text>{tool.name}: {tool.goal}</Text>
+                      </li>))}
+                    </ul>
+                  </> : <Row gap="4px">
+                    <Icon icon="BorderRadius" size="sm" color="light.700" />
+                    <Text color="light.700">{t.noToolToBeUsed}</Text>
+                  </Row>
+                  }
+                </>
+              }
+            </Column>
+          </Accordion>
+        </StyledCard>
+      </Row>
     </li>
   )
 }
 
-export const StepsList = ({ steps, chatId, messageId }: Props) => {
+export const StepsPlaceholder = () => {
+  const t = useTranslate(dictionary)
+  return <Card gap="8px">
+    <Row gap="8px">
+      <ProgressCircular colorScheme="inverse" size="xs" />
+      <Text color="light.700">{t.generatingPlan}</Text>
+    </Row>
+    <Text color="light.700">
+      {t.analyzingRequirements}
+    </Text>
+    <Row gap="12px">
+      <Skeleton height="31px" width="148px" bgLevel={600} />
+      <Skeleton height="31px" width="148px" bgLevel={600} />
+      <Skeleton height="31px" width="148px" bgLevel={600} />
+      <Skeleton height="31px" width="148px" bgLevel={600} />
+    </Row>
+  </Card>
+}
+
+const AwaitingApproval = ({ customApproveText, chatId }: { chatId: string, customApproveText?: string }) => {
+  const t = useTranslate(dictionary)
+  const chat = useChat(chatId)
+
+  const onAnswer = (response: string) => {
+    chat.pushMessage(ChatEntry.createUserEntry('', false, undefined, undefined, response))
+  }
+
+  return <>
+    <Row gap="8px">
+      <Button colorScheme="light" onClick={() => onAnswer(t.cancel)}>
+        <Row gap="8px">
+          <Icon icon="Stop" />
+          {t.cancel}
+        </Row>
+      </Button>
+
+      <Button colorScheme="inverse" onClick={() => onAnswer(customApproveText ?? t.approve)}>
+        <Row gap="8px">
+          <Icon group="fill" icon="Play" />
+          {customApproveText ?? t.approve}
+        </Row>
+      </Button>
+    </Row>
+  </>
+}
+
+export const ToolStepsList = ({ toolStep, messageId, chatId }: { toolStep: ToolChatStep, messageId: number, chatId: string }) => {
+  const t = useTranslate(dictionary)
+  const chat = useCurrentChat()
+  const messages = useCurrentChatMessages()
+  const inputParsed = `\`\`\`json
+    ${JSON.stringify(toolStep?.input, null, 2)}
+  \`\`\``
+
+  const tool = useMemo(() => {
+    if (!toolStep) return undefined
+    return toolStep.attempts?.[0].tools?.[0]
+  }, [toolStep])
+
+  useEffect(() => {
+    if (!toolStep) return undefined
+    const executionId = toolStep.attempts?.[0].tools?.[0].executionId
+    if (!executionId) return
+
+    updateToolStep(messages, executionId, toolStep.status)
+
+  }, [messages, toolStep, toolStep.status])
+
+  return <>
+    {toolStep && tool ? <AnimatedHeight>
+      <div className="steps">
+        <Badge colorPalette="yellow" appearance="square">
+          <Icon icon="StopWatch" />
+          <Text>{tool.name} {t.keepWorking}</Text>
+        </Badge>
+        <Card mt="16px" gap="8px" bgLevel={500}>
+          <Row>
+            <ImageWithFallback src={tool.image} width="32px" fallback={<IconBox appearance="circle" icon="StackSpot" />} />
+            <Text>{tool.name}</Text>
+          </Row>
+
+          <Text>
+            {toolStep.user_question}
+          </Text>
+
+          <Accordion header={expand => <Row gap="8px" mb="4px">
+            <Card p="4px" bgLevel={400}>{expand}</Card>
+            <Text > {t.viewDetails} </Text>
+          </Row>}>
+            <Markdown onCopyCode={(code) => onCopyCode(code, `${messageId}`, chat)}>
+              {inputParsed}
+            </Markdown>
+          </Accordion>
+
+          <AwaitingApproval customApproveText={t.approveTool} chatId={chatId} />
+        </Card>
+      </div>
+    </AnimatedHeight> : null}
+  </>
+}
+
+export const StepsList = ({ steps, messageId, chatId, userHasAlreadyAnswered }: Props) => {
   const t = useTranslate(dictionary)
-  const [isExpanded, setExpanded] = useState(false)
-  const actualSteps = steps.filter(s => s.type === 'step')
+
+  const filteredSteps = steps.filter(s => s.type === 'step')
+  const actualSteps = useMemo(() => filteredSteps.filter((item) => {
+    const messageIdForStep = planningToolDictionaryHelper.getMessageIdFromStepId(item.id)
+    if (!messageIdForStep) {
+      planningToolDictionaryHelper.setMessageIdForStepId(messageId, item.id)
+      return true
+    } else if (messageIdForStep === messageId) {
+      return true
+    }
+    // If a step is from a planning and it is already inserted in the planningToolDictionaryHelper it means the step is already in a 
+    // previous message and we do not want to show it again (for instance, when required a approval in the planning, we will have 
+    // two messages with the same step id one for the planning awaiting and one for the planning end, so we want to show only one) 
+    return false
+  }), [filteredSteps])
+
+  const planning = steps.filter(s => s.type === 'planning')
+
+  useEffect(() => {
+    actualSteps.map((item) => {
+      const executionId = item.attempts[0]?.tools?.[0].executionId
+      if (executionId) {
+        planningToolDictionaryHelper.setMessageIdPlanningStepToolExecutionId(`${messageId}`, executionId)
+      }
+    })
+  }, [actualSteps, messageId])
+
+  const toolsStep = findLast(steps, s => s.type === 'tool')
+
+  useEffect(() => {
+    const executionId = toolsStep?.attempts?.[0]?.tools?.[0]?.executionId
+    if (executionId) {
+      planningToolDictionaryHelper.setMessageIdToolStepToolExecutionId(`${messageId}`, executionId)
+    }
+  }, [toolsStep, messageId])
+
+  const planningGoal = planning?.[0]?.goal
+  const isLastStepDone = actualSteps[actualSteps.length - 1]?.status !== 'running' &&
+    actualSteps[actualSteps.length - 1]?.status !== 'pending'
+  const totalTools = useMemo(() => actualSteps?.reduce((sum, step) => {
+    const firstAttempt = step.attempts && step.attempts[0]
+    const toolsCount = firstAttempt.tools?.length ?? 0
+    return sum + toolsCount
+  }, 0), [steps])
+
   let currentStepIndex = findLastIndex(actualSteps, s => s.status === 'running' || s.status === 'success')
   if (currentStepIndex === -1) currentStepIndex = 0
-  const widget = useWidget()
 
-  function openToolsPanel() {
-    widget.set('currentMessageInPanel', { chatId, messageId })
-    widget.set('panel', 'steps')
-  }
-  
-  return (
-    <AnimatedHeight>
+  return (<>
+    {actualSteps.length > 0 ? <AnimatedHeight>
       <div className="steps">
-        <ul>
-          {isExpanded
-            ? actualSteps.map((s, i) => <Step step={s} key={i} index={i + 1} total={actualSteps.length} />)
-            : <Step
-              step={actualSteps[currentStepIndex]}
-              index={currentStepIndex + 1}
-              total={actualSteps.length}
-              onClick={() => setExpanded(true)}
-            />
-          }
+        <Row gap="14px" mb="8px" ml="5px">
+          <Icon icon="Target" size="sm" color="light.600" />
+          <Text>{t.executionPlan}</Text>
+        </Row>
+
+        <ul className="steps-list">
+          {actualSteps.map((s, i) => <Step step={s} key={i} index={i + 1} />)}
+
+          <Step
+            step={{ status: 'target', input: planningGoal, attempts: [] }}
+            index={currentStepIndex + 1}
+            total={actualSteps.length}
+            totalTools={totalTools}
+            isAllDone={isLastStepDone}
+          />
         </ul>
-        {isExpanded && <div className="step-actions">
-          <Button colorScheme="light" size="sm" onClick={() => setExpanded(false)}>{t.hideSteps}</Button>
-          <Button colorScheme="light" size="sm" className="icon-button details" onClick={openToolsPanel}>
-            <Icon group="fill" icon="Play" size="xs" />
-            {t.detailed}
-          </Button>
-        </div>}
+
+        <Column gap="8px" mt="8px">
+          <Divider colorScheme="light" />
+          <Text color="light.700">{planning?.[0]?.user_question}</Text>
+          {!userHasAlreadyAnswered && planning?.[0]?.status === 'awaiting_approval' &&  <AwaitingApproval chatId={chatId} />}
+        </Column>
+
       </div>
-    </AnimatedHeight>
+    </AnimatedHeight> : null}
+
+    {toolsStep && toolsStep.status === 'awaiting_approval' && !userHasAlreadyAnswered && 
+      <ToolStepsList toolStep={toolsStep} messageId={messageId} chatId={chatId} />}
+  </>
   )
 }
 
+export const ViewToolsDetails = ({ chatId }: { chatId: string }) => {
+  const t = useTranslate(dictionary)
+  const messages = useChatMessages(chatId)
+  const messageId = useMemo(() => {
+    const messageWithPlanning = findLast(messages, (message) => {
+      const steps = message.getValue().steps
+      const planningStep = steps?.find((step) => step.type === 'planning' && step.status === 'success')
+      return planningStep ? true : false
+    })
+    return messageWithPlanning?.id
+
+  }, [messages])
+  const widget = useWidget()
+
+  function openToolsPanel() {
+    if (messageId) {
+      widget.set('currentMessageInPanel', { chatId, messageId })
+      widget.set('panel', 'steps')
+    }
+  }
+
+  return <>
+    <div className="step-actions">
+      <Button colorScheme="light" size="sm" className="icon-button details" onClick={openToolsPanel}>
+        <Icon group="fill" icon="Play" size="xs" />
+        {t.detailed}
+      </Button>
+    </div>
+  </>
+}
+
 const dictionary = {
   en: {
     step: 'Step',
     hideSteps: 'Hide steps',
     detailed: 'View detailed mode',
+    generatingPlan: 'Generating execution plan...',
+    analyzingRequirements: 'Analyzing task requirements',
+    executionPlan: 'Execution Plan',
+    planGoal: 'Plan Goal',
+    toolsToBeExecuted: 'Tools to be executed',
+    totalSteps: 'Total Steps',
+    totalTools: 'Total Tools',
+    stepGoal: 'Step Goal',
+    cancel: 'Cancel execution',
+    approve: 'Approve & Execute Plan',
+    keepWorking: 'will keep working after your answer',
+    viewDetails: 'View details',
+    approveTool: 'Approve execution',
+    pendingReview: 'Pending review',
+    noToolToBeUsed: 'No tool will be needed',
   },
   pt: {
     step: 'Passo',
     hideSteps: 'Esconder passos',
     detailed: 'Ver modo detalhado',
+    generatingPlan: 'Gerando plano de execução...',
+    analyzingRequirements: 'Analisando os requisitos da task',
+    executionPlan: 'Plano de Execução',
+    planGoal: 'Finalidade do Plano',
+    toolsToBeExecuted: 'Tools a serem executadas',
+    totalSteps: 'Total de Passos',
+    totalTools: 'Total de Tools',
+    stepGoal: 'Objetivo do passo',
+    cancel: 'Cancelar execução',
+    approve: 'Aprovar & Executar plano',
+    keepWorking: 'continuará trabalhando após a sua resposta',
+    viewDetails: 'Ver detalhes',
+    approveTool: 'Aprovar execução',
+    pendingReview: 'Revisão pendente',
+    noToolToBeUsed: 'Nenhuma tool será necessária',
   },
 } satisfies Dictionary

package/src/views/Chat/styled.ts

@@ -243,20 +243,27 @@ export const ChatList: IStyledComponentBase<
   }
 
   .steps {
-    ul {
+    .steps-list {
       list-style: none;
       margin: 0;
       padding: 0;
       display: flex;
       flex-direction: column;
       gap: 6px;
+    }
 
-      li {
-        display: flex;
-        flex-direction: row;
-        gap: 4px;
-        align-items: center;
+    .tools-list {
+      list-style: disc;
+      margin: 0;
+      padding-left: 24px;
+      display: block;
+      ::marker {
+        color: ${theme.color.light.contrastText};
+      }
+    }
 
+    ul {
+      li {
         &[role="button"] {
           cursor: pointer;
         }
@@ -275,7 +282,6 @@ export const ChatList: IStyledComponentBase<
         }
 
         .step-title {
-          line-height: 0.75rem;
           overflow: hidden;
           text-overflow: ellipsis;
           display: -webkit-box;
@@ -285,8 +291,10 @@ export const ChatList: IStyledComponentBase<
       }
     }
 
-    .step-actions {
-      margin-top: 8px;
+  }
+
+  .step-actions {
+      margin-top: 16px;
       display: flex;
       gap: 6px;
 
@@ -296,7 +304,6 @@ export const ChatList: IStyledComponentBase<
         align-items: center;
       }
     }
-  }
 
   .markdown img {
     max-width: 70%;

package/src/views/MessageInput/ButtonBar.tsx

@@ -1,6 +1,7 @@
 import { IconButton, Row } from '@stack-spot/citric-react'
 import { useCurrentChat, useCurrentChatState, useWidget } from '../../context/hooks'
 import { useMessageInputDictionary } from './dictionary'
+import { ModelSwitcher } from './ModelSwitcher'
 import { SelectContent } from './SelectContent'
 import { SelectionBarWrapper } from './styled'
 
@@ -29,6 +30,7 @@ export const ButtonBar = ({ onSend, isLoading }: SelectionBarProps) => {
           <IconButton icon="Code" appearance="square" aria-label={t.code} title={t.code} onClick={() => widget.set('panel', 'editor')} />
         )}
       </Row>
+      <ModelSwitcher />
       {isLoading ? (
         <IconButton
           icon="Stop"

package/src/views/MessageInput/ModelSwitcher.tsx

@@ -0,0 +1,127 @@
+import { Button, Icon } from '@stack-spot/citric-react'
+import { SelectionList } from '@stack-spot/portal-components/SelectionList'
+import { agentToolsClient, genAiInferenceClient } from '@stack-spot/portal-network'
+import { AgentLlmModelDto, AgentModel } from '@stack-spot/portal-network/api/agent-tools'
+import { LlmModelsResponse, PaginatedResponseLlmModelsResponse } from '@stack-spot/portal-network/api/genAiInference'
+import { Dispatch, ReactElement, useState } from 'react'
+import { CSSProperties } from 'styled-components'
+import { useCurrentChat, useCurrentChatState } from '../../context/hooks'
+import { useMessageInputDictionary } from './dictionary'
+import { RowWrapperStyled } from './styled'
+
+const styles = {
+  selection: {
+    position: 'absolute',
+    bottom: 'calc(100% + 10px)',
+    right: '0',
+    margin: '0',
+    'ul': {
+      margin: '0px',
+    },
+  },
+}
+
+export const ModelSwitcher = () => {
+  const t = useMessageInputDictionary()
+  const agent = useCurrentChatState('agent')
+  const chat = useCurrentChat()
+  const [visibleMenu, setVisibleMenu] = useState(false)
+  const [agentData] = agentToolsClient.agent.useStatefulQuery({ agentId: agent?.id || '' })
+  const [selectedLLMName, setSelectedLLMName] = useState<string | undefined>()
+  const [models] = genAiInferenceClient.listModels.useStatefulQuery({ pageSize: 999 })
+
+  const { modelName, listItems } = getModelData(chat, setSelectedLLMName, setVisibleMenu, selectedLLMName, agentData, models)
+
+  return (
+    <RowWrapperStyled>
+      <Button
+        className="button-select-model"
+        colorScheme="light"
+        size="sm"
+        aria-label={t.agent}
+        title={t.agent}
+        onClick={() => setVisibleMenu(state => !state)}
+      >
+        {modelName}
+        <Icon icon="ChevronDown" group="fill" size="sm" />
+      </Button>
+      <SelectionList
+        id="menuModelSwitcher"
+        visible={visibleMenu}
+        onHide={() => setVisibleMenu(false)}
+        items={listItems || []}
+        style={styles.selection as CSSProperties}
+      />
+    </RowWrapperStyled>
+  )
+}
+
+
+function getModelData(
+  chat: any, 
+  setSelectedModelName: Dispatch<React.SetStateAction<string | undefined>>, 
+  setVisibleMenu: Dispatch<React.SetStateAction<boolean>>, 
+  selectedModelName?: string,
+  agent?: AgentModel,
+  models?: PaginatedResponseLlmModelsResponse) {
+
+  const modelListData = parseModelList(
+    setSelectedModelName, setVisibleMenu, chat,
+    agent?.visibility_level !== 'built_in' && !!agent?.model_id ? agent?.available_llm_models : models?.items,
+  )
+
+  if (agent?.visibility_level === 'built_in' || !agent?.model_id) {
+    const modelDefaultFromList = modelListData.find((model) => model.active)
+    return { modelName: modelDefaultFromList?.label, listItems: modelListData }
+  }
+
+  const listItems = 
+    modelListData && modelListData?.length > 0 ? modelListData : 
+      [{ 
+        active: true, 
+        label: agent?.available_llm_models?.find((model) => model.is_default)?.model_name || agent?.model_name || '', 
+        icon: <Icon icon="StackSpot" />, 
+      }]
+
+  const modelSelectedName = selectedModelName || agent?.model_name ||
+    agent?.available_llm_models?.find((model) => model.is_default)?.model_name
+
+  return { modelName: modelSelectedName, listItems }
+}
+
+function parseModelList(
+  setSelectedModelName: Dispatch<React.SetStateAction<string | undefined>>, 
+  setVisibleMenu: Dispatch<React.SetStateAction<boolean>>, 
+  chat: any,
+  listModel?: Array<AgentLlmModelDto | LlmModelsResponse>, 
+) {
+  const data = Array<{active?: boolean, label: string, icon: ReactElement, onClick: VoidFunction  }>()
+
+  listModel?.forEach((model) => {
+    if ('model_id' in model) {
+      data.push({
+        active: chat.get('selected_model_id') ? chat.get('selected_model_id') === model.model_id : model.is_default,
+        label: model?.model_name || 'LLM',
+        icon: <Icon icon="Bookmark"/>,
+        onClick: () => {
+          setSelectedModelName(model.model_name)
+          chat.set('selected_model_id', model.model_id)
+          setVisibleMenu(false)
+        },
+      })
+    }
+    else {
+      data.push({
+        active: chat.get('selected_model_id') ? chat.get('selected_model_id') === model.id : model.resources[0]?.is_default,
+        label: model?.display_name || 'LLM',
+        icon: <Icon icon="Bookmark"/>,
+        onClick: () => {
+          setSelectedModelName(model.display_name)
+          chat.set('selected_model_id', model.id)
+          setVisibleMenu(false)
+        },
+      })
+    }
+  })
+  return data
+}

package/src/views/MessageInput/styled.ts

@@ -1,3 +1,4 @@
+import { Flex } from '@citric/core'
 import { theme } from '@stack-spot/portal-theme'
 import { styled } from 'styled-components'
 
@@ -302,3 +303,14 @@ export const MessageInputBox = styled.div`
     }
   }
 `
+export const RowWrapperStyled = styled(Flex)`
+  width: 100%;
+  justify-content: end;
+  margin-right: 4px;
+  ul {
+    margin: 0;
+  }
+  .button-select-model {
+    border-radius: 15px !important;
+  }
+`

package/src/views/Steps/FlowChart/NodeStep.tsx

@@ -37,7 +37,7 @@ export const NodeStep = ({ data: { step, index, nextStatus, onClick } }: Props)
         {!!step.attempts[0].tools?.length && <StackedBadge
           label={t.tools}
           images={step.attempts[0].tools?.slice(0, 3).map(
-            tool => ({ key: tool.id, name: tool.name, url: tool.image, icon: <Icon icon="Cog" /> }),
+            tool => ({ key: tool.id, name: tool.name ?? '', url: tool.image, icon: <Icon icon="Cog" /> }),
           )}
         />}
       </div>}

package/src/views/Steps/FlowChart/layout.ts

@@ -8,6 +8,7 @@ const nodesSizes = {
   step: stepNodeSize,
   planning: planningNodeSize,
   answer: answerNodeSize,
+  tool: stepNodeSize,
 }
 
 export function useLayoutedElements(nodes: NodeWithoutLayout[], edges: Edge[]) {