diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message.tsx b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message.tsx
index 555461de16d77..7e22c2bbe5db8 100644
--- a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message.tsx
+++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message.tsx
@@ -88,7 +88,7 @@ export function WelcomeMessage({
             connectors={connectors}
             onSetupConnectorClick={handleConnectorClick}
           />
-          {knowledgeBase.status.value?.enabled ? (
+          {knowledgeBase.status.value?.enabled && connectors.connectors?.length ? (
             <WelcomeMessageKnowledgeBase connectors={connectors} knowledgeBase={knowledgeBase} />
           ) : null}
         </EuiFlexItem>
diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message_knowledge_base.test.tsx b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message_knowledge_base.test.tsx
new file mode 100644
index 0000000000000..656d0d3a867e3
--- /dev/null
+++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message_knowledge_base.test.tsx
@@ -0,0 +1,266 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React from 'react';
+import { act, render, screen } from '@testing-library/react';
+
+import { WelcomeMessageKnowledgeBase } from './welcome_message_knowledge_base';
+import type { UseGenAIConnectorsResult } from '../hooks/use_genai_connectors';
+import type { UseKnowledgeBaseResult } from '../hooks/use_knowledge_base';
+
+describe('WelcomeMessageKnowledgeBase', () => {
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+  type StatusType = NonNullable<UseKnowledgeBaseResult['status']['value']>;
+  type EndpointType = StatusType['endpoint'];
+  const endpoint: EndpointType = {
+    inference_id: 'obs_ai_assistant_kb_inference',
+    task_type: 'sparse_embedding',
+    service: 'elasticsearch',
+    service_settings: {
+      num_threads: 1,
+      model_id: '.elser_model_2',
+      adaptive_allocations: {
+        enabled: true,
+        min_number_of_allocations: 1,
+      },
+    },
+  };
+  const initKnowledgeBase: UseKnowledgeBaseResult = {
+    isInstalling: false,
+    install: jest.fn(),
+    installError: undefined,
+    status: {
+      value: {
+        ready: false,
+        enabled: true,
+        errorMessage: 'error',
+      },
+      loading: false,
+      refresh: jest.fn(),
+    },
+  };
+  const defaultConnectors: UseGenAIConnectorsResult = {
+    connectors: [
+      {
+        id: 'default-connector-id',
+        actionTypeId: 'action-type-id',
+        name: 'Default Connector',
+        isPreconfigured: false,
+        isDeprecated: false,
+        isSystemAction: false,
+        referencedByCount: 0,
+      },
+    ],
+    selectedConnector: undefined,
+    loading: false,
+    error: undefined,
+    selectConnector: jest.fn(),
+    reloadConnectors: jest.fn(),
+  };
+  function renderComponent({
+    knowledgeBase,
+    connectors,
+  }: {
+    knowledgeBase: Partial<UseKnowledgeBaseResult>;
+    connectors: Partial<UseGenAIConnectorsResult>;
+  }) {
+    const mergedKnowledgeBase: UseKnowledgeBaseResult = {
+      ...initKnowledgeBase,
+      ...knowledgeBase,
+      status: {
+        ...initKnowledgeBase.status,
+        ...knowledgeBase.status,
+      },
+    };
+
+    return render(
+      <WelcomeMessageKnowledgeBase
+        knowledgeBase={mergedKnowledgeBase}
+        connectors={defaultConnectors}
+      />
+    );
+  }
+
+  it('renders "Setting up Knowledge base" message while inference endpoint is installing', () => {
+    renderComponent({
+      knowledgeBase: {
+        isInstalling: true,
+      },
+      connectors: defaultConnectors,
+    });
+
+    expect(
+      screen.getByText('We are setting up your knowledge base', { exact: false })
+    ).toBeInTheDocument();
+
+    expect(screen.getByText('Setting up Knowledge base', { exact: false })).toBeInTheDocument();
+  });
+  it('renders "Setting up Knowledge base" message while model is being deployed without deployment or allocation state yet being reported', () => {
+    renderComponent({
+      knowledgeBase: {
+        isInstalling: false,
+        status: {
+          value: {
+            endpoint,
+            ready: false,
+            enabled: true,
+            model_stats: { allocation_count: 0 },
+          },
+          loading: false,
+          refresh: jest.fn(),
+        },
+      },
+      connectors: defaultConnectors,
+    });
+    expect(
+      screen.getByText('We are setting up your knowledge base', { exact: false })
+    ).toBeInTheDocument();
+  });
+  it('renders "Setting up Knowledge base" message while model is being deployed and starting', () => {
+    renderComponent({
+      knowledgeBase: {
+        isInstalling: false,
+        status: {
+          value: {
+            endpoint,
+            ready: false,
+            enabled: true,
+            model_stats: {
+              deployment_state: 'starting',
+              allocation_state: 'starting',
+            },
+          },
+          loading: false,
+          refresh: jest.fn(),
+        },
+      },
+      connectors: defaultConnectors,
+    });
+
+    expect(
+      screen.getByText('We are setting up your knowledge base', { exact: false })
+    ).toBeInTheDocument();
+  });
+  it('displays success message after installation and hides it after timeout', async () => {
+    jest.useFakeTimers();
+
+    // Step 1: Initially not installed
+    const { rerender } = renderComponent({
+      knowledgeBase: {
+        isInstalling: true,
+      },
+      connectors: defaultConnectors,
+    });
+
+    // Step 2: Now it's ready
+    await act(async () => {
+      rerender(
+        <WelcomeMessageKnowledgeBase
+          knowledgeBase={{
+            ...initKnowledgeBase,
+            status: {
+              ...initKnowledgeBase.status,
+              value: {
+                ready: true,
+                enabled: true,
+                model_stats: { deployment_state: 'started', allocation_state: 'started' },
+              },
+            },
+          }}
+          connectors={defaultConnectors}
+        />
+      );
+    });
+
+    // the success message should appear
+    expect(screen.queryByText(/Knowledge base successfully installed/i)).toBeInTheDocument();
+
+    // fast-forward until the success message would disappear
+    await act(async () => {
+      jest.runOnlyPendingTimers();
+    });
+
+    // now it should be gone
+    expect(screen.queryByText('Knowledge base successfully installed')).toBeNull();
+  });
+
+  it('renders no install messages when model has been installed and ready', () => {
+    // component should render nothing in this state (null)
+    renderComponent({
+      knowledgeBase: {
+        isInstalling: false,
+        status: {
+          ...initKnowledgeBase.status,
+          value: {
+            ready: true,
+            enabled: true,
+            model_stats: { deployment_state: 'started', allocation_state: 'started' },
+          },
+        },
+      },
+      connectors: defaultConnectors,
+    });
+    expect(screen.queryByText(/We are setting up your knowledge base/i)).not.toBeInTheDocument();
+    expect(screen.queryByText(/Your Knowledge base hasn't been set up./i)).not.toBeInTheDocument();
+  });
+
+  it('renders knowledge base install and model state inspect when not installing and the inference endpoint installation has an error', () => {
+    renderComponent({
+      knowledgeBase: {
+        isInstalling: false,
+        installError: new Error('inference endpoint failed to install'),
+      },
+      connectors: defaultConnectors,
+    });
+    expect(
+      screen.getByText("Your Knowledge base hasn't been set up", { exact: false })
+    ).toBeInTheDocument();
+    expect(screen.getByText('Install Knowledge base', { exact: false })).toBeInTheDocument();
+    expect(screen.getByText('Inspect issues', { exact: false })).toBeInTheDocument();
+  });
+
+  it('renders knowledge base install and model state inspect when not installing and no errors', () => {
+    // this can happen when you have a preconfigured connector,
+    // we don't automatically install in this case and just show the same UI as if there was an issue/error
+    renderComponent({
+      knowledgeBase: {
+        isInstalling: false,
+      },
+      connectors: defaultConnectors,
+    });
+    expect(
+      screen.getByText("Your Knowledge base hasn't been set up", { exact: false })
+    ).toBeInTheDocument();
+    expect(screen.getByText('Install Knowledge base', { exact: false })).toBeInTheDocument();
+    expect(screen.getByText('Inspect issues', { exact: false })).toBeInTheDocument();
+  });
+
+  it('renders knowledge base install and model state inspect when not installing and model is not ready', () => {
+    // this state can occur if the model is having a deployment problem
+    renderComponent({
+      knowledgeBase: {
+        isInstalling: false,
+        status: {
+          ...initKnowledgeBase.status,
+          value: {
+            ready: false,
+            enabled: true,
+            model_stats: { deployment_state: 'failed', allocation_state: 'started' },
+          },
+        },
+      },
+      connectors: defaultConnectors,
+    });
+    expect(
+      screen.getByText("Your Knowledge base hasn't been set up", { exact: false })
+    ).toBeInTheDocument();
+    expect(screen.getByText('Install Knowledge base', { exact: false })).toBeInTheDocument();
+    expect(screen.getByText('Inspect issues', { exact: false })).toBeInTheDocument();
+  });
+});
diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message_knowledge_base.tsx b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message_knowledge_base.tsx
index 72653473c41ae..8913d32864288 100644
--- a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message_knowledge_base.tsx
+++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/welcome_message_knowledge_base.tsx
@@ -40,44 +40,61 @@ export function WelcomeMessageKnowledgeBase({
   const [isPopoverOpen, setIsPopoverOpen] = useState(false);
   const handleClosePopover = () => setIsPopoverOpen(false);
 
-  const [checkForInstallStatus, setCheckForInstallStatus] = useState(false);
+  const [pollKnowledgeBaseStatus, setPollKnowledgeBaseStatus] = useState(false);
 
-  // When the knowledge base is installed, show a success message for 3 seconds
+  // Tracks whether the inference endpoint creation process has started
+  const inferenceEndpointIsInstalling = knowledgeBase.isInstalling;
+
+  // Tracks whether the model is fully ready
+  const modelIsReady = knowledgeBase.status.value?.ready === true;
+
+  // Determines if the model deployment is still in progress
+  // This happens when the model is not ready but the endpoint exists
+  const modelDeploymentInProgress = !modelIsReady && !!knowledgeBase.status.value?.endpoint;
+
+  // Determines if the overall installation process is ongoing
+  // Covers both the endpoint setup phase and the model deployment phase
+  const isInstalling = inferenceEndpointIsInstalling || modelDeploymentInProgress;
+  // start polling kb status if inference endpoint is being created or has been created but model isn't ready
+  useEffect(() => {
+    if (isInstalling) {
+      setPollKnowledgeBaseStatus(true);
+    }
+  }, [isInstalling]);
+
+  // When the knowledge base is installed and ready, show a success message for 3 seconds
   useEffect(() => {
-    if (previouslyNotInstalled && knowledgeBase.status.value?.ready) {
+    if (previouslyNotInstalled && modelIsReady) {
       setTimeoutTime(3000);
       reset();
       setShowHasBeenInstalled(true);
     }
-  }, [knowledgeBase.status.value?.ready, previouslyNotInstalled, reset]);
+  }, [modelIsReady, previouslyNotInstalled, reset]);
 
-  // When the knowledge base is installed, stop checking for install status
+  // When the knowledge base is ready, stop polling for status
   useEffect(() => {
-    if (!checkForInstallStatus && knowledgeBase.status.value?.ready) {
-      setCheckForInstallStatus(false);
+    if (modelIsReady) {
+      setPollKnowledgeBaseStatus(false);
     }
-  }, [checkForInstallStatus, knowledgeBase.status.value?.ready]);
+  }, [modelIsReady]);
 
-  // Check for install status every 5 seconds
+  // poll for knowledge base status every 5 seconds
   useInterval(
     () => {
       knowledgeBase.status.refresh();
     },
-    checkForInstallStatus ? 5000 : null
+    pollKnowledgeBaseStatus ? 5000 : null
   );
 
-  const handleRetryInstall = async () => {
-    setCheckForInstallStatus(true);
+  // gets called if there was an error previously during install or user has a preconfigured connector
+  // and is first time installing
+  const handleInstall = async () => {
     setIsPopoverOpen(false);
-
-    await knowledgeBase.install().then(() => {
-      setCheckForInstallStatus(false);
-    });
+    await knowledgeBase.install();
   };
-
-  return knowledgeBase.status.value?.ready !== undefined ? (
+  return modelIsReady !== undefined ? (
     <>
-      {knowledgeBase.isInstalling ? (
+      {isInstalling ? (
         <>
           <EuiText color="subdued" size="s">
             {i18n.translate('xpack.aiAssistant.welcomeMessage.weAreSettingUpTextLabel', {
@@ -100,69 +117,72 @@ export function WelcomeMessageKnowledgeBase({
         </>
       ) : null}
 
-      {connectors.connectors?.length ? (
-        (!knowledgeBase.isInstalling && knowledgeBase.installError) ||
-        (!knowledgeBase.isInstalling &&
-          knowledgeBase.status.loading === false &&
-          knowledgeBase.status.value.ready === false) ? (
-          <>
-            <EuiText color="subdued" size="s">
-              {i18n.translate(
-                'xpack.aiAssistant.welcomeMessageKnowledgeBase.yourKnowledgeBaseIsNotSetUpCorrectlyLabel',
-                { defaultMessage: `Your Knowledge base hasn't been set up.` }
-              )}
-            </EuiText>
-
-            <EuiSpacer size="m" />
-
-            <EuiFlexGroup justifyContent="center">
-              <EuiFlexItem grow={false}>
-                <div>
-                  <EuiButton
-                    color="primary"
-                    data-test-subj="observabilityAiAssistantWelcomeMessageSetUpKnowledgeBaseButton"
-                    fill
-                    isLoading={checkForInstallStatus}
-                    iconType="importAction"
-                    onClick={handleRetryInstall}
-                  >
-                    {i18n.translate('xpack.aiAssistant.welcomeMessage.retryButtonLabel', {
-                      defaultMessage: 'Install Knowledge base',
-                    })}
-                  </EuiButton>
-                </div>
-              </EuiFlexItem>
-
-              <EuiFlexItem grow={false}>
-                <EuiPopover
-                  button={
-                    <EuiButtonEmpty
-                      data-test-subj="observabilityAiAssistantWelcomeMessageInspectErrorsButton"
-                      iconType="inspect"
-                      onClick={() => setIsPopoverOpen(!isPopoverOpen)}
+      {
+        // not currently installing
+        // and has an inference install error (timeout, etc) or model is not ready
+        // this state is when the user has a preconfigured connector and we prompt to install
+        // or there was a problem deploying the model
+        !isInstalling ? (
+          knowledgeBase.installError || !modelIsReady ? (
+            <>
+              <EuiText color="subdued" size="s">
+                {i18n.translate(
+                  'xpack.aiAssistant.welcomeMessageKnowledgeBase.yourKnowledgeBaseIsNotSetUpCorrectlyLabel',
+                  { defaultMessage: `Your Knowledge base hasn't been set up.` }
+                )}
+              </EuiText>
+
+              <EuiSpacer size="m" />
+
+              <EuiFlexGroup justifyContent="center">
+                <EuiFlexItem grow={false}>
+                  <div>
+                    <EuiButton
+                      color="primary"
+                      data-test-subj="observabilityAiAssistantWelcomeMessageSetUpKnowledgeBaseButton"
+                      fill
+                      isLoading={false}
+                      iconType="importAction"
+                      onClick={handleInstall}
                     >
-                      {i18n.translate(
-                        'xpack.aiAssistant.welcomeMessage.inspectErrorsButtonEmptyLabel',
-                        { defaultMessage: 'Inspect issues' }
-                      )}
-                    </EuiButtonEmpty>
-                  }
-                  isOpen={isPopoverOpen}
-                  panelPaddingSize="none"
-                  closePopover={handleClosePopover}
-                >
-                  <WelcomeMessageKnowledgeBaseSetupErrorPanel
-                    knowledgeBase={knowledgeBase}
-                    onRetryInstall={handleRetryInstall}
-                  />
-                </EuiPopover>
-              </EuiFlexItem>
-            </EuiFlexGroup>
-
-            <EuiSpacer size="m" />
-          </>
+                      {i18n.translate('xpack.aiAssistant.welcomeMessage.retryButtonLabel', {
+                        defaultMessage: 'Install Knowledge base',
+                      })}
+                    </EuiButton>
+                  </div>
+                </EuiFlexItem>
+
+                <EuiFlexItem grow={false}>
+                  <EuiPopover
+                    button={
+                      <EuiButtonEmpty
+                        data-test-subj="observabilityAiAssistantWelcomeMessageInspectErrorsButton"
+                        iconType="inspect"
+                        onClick={() => setIsPopoverOpen(!isPopoverOpen)}
+                      >
+                        {i18n.translate(
+                          'xpack.aiAssistant.welcomeMessage.inspectErrorsButtonEmptyLabel',
+                          { defaultMessage: 'Inspect issues' }
+                        )}
+                      </EuiButtonEmpty>
+                    }
+                    isOpen={isPopoverOpen}
+                    panelPaddingSize="none"
+                    closePopover={handleClosePopover}
+                  >
+                    <WelcomeMessageKnowledgeBaseSetupErrorPanel
+                      knowledgeBase={knowledgeBase}
+                      onRetryInstall={handleInstall}
+                    />
+                  </EuiPopover>
+                </EuiFlexItem>
+              </EuiFlexGroup>
+
+              <EuiSpacer size="m" />
+            </>
+          ) : null
         ) : null
-      ) : null}
+      }
 
       {showHasBeenInstalled ? (
         <div>
diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/hooks/use_knowledge_base.tsx b/x-pack/platform/packages/shared/kbn-ai-assistant/src/hooks/use_knowledge_base.tsx
index a9feb325e6d6c..5ab4193d6ad48 100644
--- a/x-pack/platform/packages/shared/kbn-ai-assistant/src/hooks/use_knowledge_base.tsx
+++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/hooks/use_knowledge_base.tsx
@@ -62,8 +62,8 @@ export function useKnowledgeBase(): UseKnowledgeBaseResult {
           }
           setInstallError(error);
           notifications!.toasts.addError(error, {
-            title: i18n.translate('xpack.aiAssistant.errorSettingUpKnowledgeBase', {
-              defaultMessage: 'Could not set up Knowledge Base',
+            title: i18n.translate('xpack.aiAssistant.errorSettingUpInferenceEndpoint', {
+              defaultMessage: 'Could not create inference endpoint',
             }),
           });
         })
diff --git a/x-pack/platform/plugins/private/translations/translations/fr-FR.json b/x-pack/platform/plugins/private/translations/translations/fr-FR.json
index 49d99cf45d346..d51588efc972e 100644
--- a/x-pack/platform/plugins/private/translations/translations/fr-FR.json
+++ b/x-pack/platform/plugins/private/translations/translations/fr-FR.json
@@ -9793,7 +9793,6 @@
     "xpack.aiAssistant.couldNotFindConversationTitle": "Conversation introuvable",
     "xpack.aiAssistant.disclaimer.disclaimerLabel": "Ce chat est soutenu par une intégration avec votre fournisseur LLM. Il arrive que les grands modèles de langage (LLM) présentent comme correctes des informations incorrectes. Elastic prend en charge la configuration ainsi que la connexion au fournisseur LLM et à votre base de connaissances, mais n'est pas responsable des réponses fournies par le LLM.",
     "xpack.aiAssistant.emptyConversationTitle": "Nouvelle conversation",
-    "xpack.aiAssistant.errorSettingUpKnowledgeBase": "Impossible de configurer la base de connaissances",
     "xpack.aiAssistant.errorUpdatingConversation": "Impossible de mettre à jour la conversation",
     "xpack.aiAssistant.executedFunctionFailureEvent": "impossible d'exécuter la fonction {functionName}",
     "xpack.aiAssistant.failedToGetStatus": "Échec de l'obtention du statut du modèle.",
diff --git a/x-pack/platform/plugins/private/translations/translations/ja-JP.json b/x-pack/platform/plugins/private/translations/translations/ja-JP.json
index 563b5f13dd00b..a7a455d45def3 100644
--- a/x-pack/platform/plugins/private/translations/translations/ja-JP.json
+++ b/x-pack/platform/plugins/private/translations/translations/ja-JP.json
@@ -9669,7 +9669,6 @@
     "xpack.aiAssistant.couldNotFindConversationTitle": "会話が見つかりません",
     "xpack.aiAssistant.disclaimer.disclaimerLabel": "このチャットは、LLMプロバイダーとの統合によって提供されています。LLMは、正しくない情報を正しい情報であるかのように表示する場合があることが知られています。Elasticは、構成やLLMプロバイダーへの接続、お客様のナレッジベースへの接続はサポートしますが、LLMの応答については責任を負いません。",
     "xpack.aiAssistant.emptyConversationTitle": "新しい会話",
-    "xpack.aiAssistant.errorSettingUpKnowledgeBase": "ナレッジベースをセットアップできませんでした",
     "xpack.aiAssistant.errorUpdatingConversation": "会話を更新できませんでした",
     "xpack.aiAssistant.executedFunctionFailureEvent": "関数{functionName}の実行に失敗しました",
     "xpack.aiAssistant.failedToGetStatus": "モデルステータスを取得できませんでした。",
diff --git a/x-pack/platform/plugins/private/translations/translations/zh-CN.json b/x-pack/platform/plugins/private/translations/translations/zh-CN.json
index 42d027dcacfb4..1da3528eebbdd 100644
--- a/x-pack/platform/plugins/private/translations/translations/zh-CN.json
+++ b/x-pack/platform/plugins/private/translations/translations/zh-CN.json
@@ -9511,7 +9511,6 @@
     "xpack.aiAssistant.couldNotFindConversationTitle": "未找到对话",
     "xpack.aiAssistant.disclaimer.disclaimerLabel": "通过集成 LLM 提供商来支持此聊天。众所周知，LLM 有时会提供错误信息，好像它是正确的。Elastic 支持配置并连接到 LLM 提供商和知识库，但不对 LLM 响应负责。",
     "xpack.aiAssistant.emptyConversationTitle": "新对话",
-    "xpack.aiAssistant.errorSettingUpKnowledgeBase": "无法设置知识库",
     "xpack.aiAssistant.errorUpdatingConversation": "无法更新对话",
     "xpack.aiAssistant.executedFunctionFailureEvent": "无法执行函数 {functionName}",
     "xpack.aiAssistant.failedToGetStatus": "无法获取模型状态。",
diff --git a/x-pack/platform/plugins/shared/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts b/x-pack/platform/plugins/shared/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts
index 6aa63c177c746..ddc644c4b6d7c 100644
--- a/x-pack/platform/plugins/shared/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts
+++ b/x-pack/platform/plugins/shared/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts
@@ -13,6 +13,7 @@ import {
   InferenceInferenceEndpointInfo,
   MlDeploymentAllocationState,
   MlDeploymentAssignmentState,
+  MlTrainedModelDeploymentAllocationStatus,
 } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import moment from 'moment';
 import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
@@ -34,8 +35,9 @@ const getKnowledgeBaseStatus = createObservabilityAIAssistantServerRoute({
     enabled: boolean;
     endpoint?: Partial<InferenceInferenceEndpointInfo>;
     model_stats?: {
-      deployment_state: MlDeploymentAssignmentState | undefined;
-      allocation_state: MlDeploymentAllocationState | undefined;
+      deployment_state?: MlDeploymentAssignmentState;
+      allocation_state?: MlDeploymentAllocationState;
+      allocation_count?: MlTrainedModelDeploymentAllocationStatus['allocation_count'];
     };
   }> => {
     const client = await service.getClient({ request });