From 6c0387d170a2caab45393222bcc32e769a2c5124 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Tue, 3 Mar 2026 19:31:55 +0530
Subject: [PATCH] Add support for Attaching knowledge base to model via UI

---
 .../src/components/add_model/AddModelForm.tsx |  1 +
 .../add_model/advanced_settings.test.tsx      |  3 +
 .../add_model/advanced_settings.tsx           | 30 +++++++++
 .../src/components/model_info_view.tsx        | 61 +++++++++++++++++++
 4 files changed, 95 insertions(+)

diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx
index 59ac63cffe6..2b3f23a35ae 100644
--- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx
@@ -358,6 +358,7 @@ const AddModelForm: React.FC<AddModelFormProps> = ({
                   teams={teams}
                   guardrailsList={guardrailsList || []}
                   tagsList={tagsList || {}}
+                  accessToken={accessToken || ""}
                 />
               </>
             )}
diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx
index 6515c67c292..9fe36e13998 100644
--- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx
@@ -13,6 +13,7 @@ describe("AdvancedSettings", () => {
         setShowAdvancedSettings={() => {}}
         guardrailsList={[]}
         tagsList={{}}
+        accessToken="test-token"
       />,
     );
   });
@@ -24,6 +25,7 @@ describe("AdvancedSettings", () => {
         setShowAdvancedSettings={() => {}}
         guardrailsList={[]}
         tagsList={{}}
+        accessToken="test-token"
       />,
     );
     fireEvent.click(getByText("Advanced Settings"));
@@ -39,6 +41,7 @@ describe("AdvancedSettings", () => {
         setShowAdvancedSettings={() => {}}
         guardrailsList={[]}
         tagsList={{}}
+        accessToken="test-token"
       />,
     );
     act(() => {
diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
index c9f5ef8a4b1..8ae90c1cbbc 100644
--- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
@@ -6,6 +6,7 @@ import TextArea from "antd/es/input/TextArea";
 import { InfoCircleOutlined } from "@ant-design/icons";
 import { Team } from "../key_team_helpers/key_list";
 import CacheControlSettings from "./cache_control_settings";
+import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
 import { Tag } from "../tag_management/types";
 import { formItemValidateJSON } from "../../utils/textUtils";
 const { Link } = Typography;
@@ -16,6 +17,7 @@ interface AdvancedSettingsProps {
   teams?: Team[] | null;
   guardrailsList: string[];
   tagsList: Record<string, Tag>;
+  accessToken: string;
 }
 
 const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
@@ -24,6 +26,7 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
   teams,
   guardrailsList,
   tagsList,
+  accessToken,
 }) => {
   const [form] = Form.useForm();
   const [customPricing, setCustomPricing] = React.useState(false);
@@ -109,6 +112,33 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
               <Switch onChange={handleCustomPricingChange} className="bg-gray-600" />
             </Form.Item>
 
+            <Form.Item
+              label={
+                <span>
+                  Attached Knowledge Bases (RAG){" "}
+                  <Tooltip title="Vector stores to use for RAG. Every request to this model will automatically retrieve context from these knowledge bases.">
+                    <a
+                      href="https://docs.litellm.ai/docs/completion/knowledgebase"
+                      target="_blank"
+                      rel="noopener noreferrer"
+                      onClick={(e) => e.stopPropagation()}
+                    >
+                      <InfoCircleOutlined style={{ marginLeft: "4px" }} />
+                    </a>
+                  </Tooltip>
+                </span>
+              }
+              name="vector_store_ids"
+              className="mt-4"
+              help="Select vector stores to attach. Requests to this model will automatically use these for RAG. Set up vector stores in Tools > Vector Stores."
+            >
+              <VectorStoreSelector
+                onChange={() => {}}
+                accessToken={accessToken}
+                placeholder="Select knowledge bases (optional)"
+              />
+            </Form.Item>
+
             <Form.Item
               label={
                 <span>
diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx
index e2fc8caa21c..40c1a3a386a 100644
--- a/ui/litellm-dashboard/src/components/model_info_view.tsx
+++ b/ui/litellm-dashboard/src/components/model_info_view.tsx
@@ -17,6 +17,7 @@ import {
   Button as TremorButton,
 } from "@tremor/react";
 import { Button, Form, Input, Modal, Select, Tooltip } from "antd";
+import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
 import { CheckIcon, CopyIcon } from "lucide-react";
 import { useEffect, useMemo, useState } from "react";
 import { copyToClipboard as utilCopyToClipboard } from "../utils/dataUtils";
@@ -245,6 +246,11 @@ export default function ModelInfoView({
       if (values.guardrails) {
         updatedLitellmParams.guardrails = values.guardrails;
       }
+      if (values.vector_store_ids !== undefined) {
+        updatedLitellmParams.vector_store_ids = Array.isArray(values.vector_store_ids)
+          ? values.vector_store_ids
+          : [];
+      }
 
       // Handle cache control settings
       if (values.cache_control && values.cache_control_injection_points?.length > 0) {
@@ -606,6 +612,9 @@ export default function ModelInfoView({
                     guardrails: Array.isArray(localModelData.litellm_params?.guardrails)
                       ? localModelData.litellm_params.guardrails
                       : [],
+                    vector_store_ids: Array.isArray(localModelData.litellm_params?.vector_store_ids)
+                      ? localModelData.litellm_params.vector_store_ids
+                      : [],
                     tags: Array.isArray(localModelData.litellm_params?.tags) ? localModelData.litellm_params.tags : [],
                     health_check_model: isWildcardModel ? localModelData.model_info?.health_check_model : null,
                     litellm_extra_params: JSON.stringify(localModelData.litellm_params || {}, null, 2),
@@ -883,6 +892,58 @@ export default function ModelInfoView({
                         )}
                       </div>
 
+                      <div>
+                        <Text className="font-medium">
+                          Attached Knowledge Bases (RAG)
+                          <Tooltip title="Vector stores used for RAG. Every request to this model will automatically retrieve context from these knowledge bases.">
+                            <a
+                              href="https://docs.litellm.ai/docs/completion/knowledgebase"
+                              target="_blank"
+                              rel="noopener noreferrer"
+                              onClick={(e) => e.stopPropagation()}
+                            >
+                              <InfoCircleOutlined style={{ marginLeft: "4px" }} />
+                            </a>
+                          </Tooltip>
+                        </Text>
+                        {isEditing ? (
+                          <Form.Item name="vector_store_ids" className="mb-0">
+                            <VectorStoreSelector
+                              onChange={() => {}}
+                              accessToken={accessToken || ""}
+                              placeholder="Select knowledge bases (optional)"
+                            />
+                          </Form.Item>
+                        ) : (
+                          <div className="mt-1 p-2 bg-gray-50 rounded">
+                            {localModelData.litellm_params?.vector_store_ids ? (
+                              Array.isArray(localModelData.litellm_params.vector_store_ids) ? (
+                                localModelData.litellm_params.vector_store_ids.length > 0 ? (
+                                  <div className="flex flex-wrap gap-1">
+                                    {localModelData.litellm_params.vector_store_ids.map(
+                                      (vsId: string, index: number) => (
+                                        <span
+                                          key={index}
+                                          className="inline-flex items-center px-2 py-1 rounded-full text-xs font-medium bg-blue-100 text-blue-800"
+                                        >
+                                          {vsId}
+                                        </span>
+                                      )
+                                    )}
+                                  </div>
+                                ) : (
+                                  "No knowledge bases attached"
+                                )
+                              ) : (
+                                String(localModelData.litellm_params.vector_store_ids)
+                              )
+                            ) : (
+                              "Not Set"
+                            )}
+                          </div>
+                        )}
+                      </div>
+
                       <div>
                         <Text className="font-medium">Tags</Text>
                         {isEditing ? (