Add support for Attaching knowledge base to model via UI

This commit is contained in:
Sameer Kankute 2026-03-03 19:31:55 +05:30
parent 67f90254ed
commit 6c0387d170
4 changed files with 95 additions and 0 deletions

View file

@ -358,6 +358,7 @@ const AddModelForm: React.FC<AddModelFormProps> = ({
teams={teams}
guardrailsList={guardrailsList || []}
tagsList={tagsList || {}}
accessToken={accessToken || ""}
/>
</>
)}

View file

@ -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(() => {

View file

@ -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>

View file

@ -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 ? (