diff --git a/docs/my-website/docs/completion/knowledgebase.md b/docs/my-website/docs/completion/knowledgebase.md index d527066e16b..e810935a37c 100644 --- a/docs/my-website/docs/completion/knowledgebase.md +++ b/docs/my-website/docs/completion/knowledgebase.md @@ -1,10 +1,10 @@ -# Using Knowledge Bases with LiteLLM +# Using Vector Stores (Knowledge Bases) with LiteLLM LiteLLM integrates with AWS Bedrock Knowledge Bases, allowing your models to access your organization's data for more accurate and contextually relevant responses. ## Quick Start -In order to use a Bedrock Knowledge Base with LiteLLM, you need to pass `knowledge_bases` as a parameter to the completion request. Where `knowledge_bases` is a list of Bedrock Knowledge Base IDs. +In order to use a Bedrock Knowledge Base with LiteLLM, you need to pass `vector_store_ids` as a parameter to the completion request. Where `vector_store_ids` is a list of Bedrock Knowledge Base IDs. ### LiteLLM Python SDK @@ -13,11 +13,11 @@ import os import litellm -# Make a completion request with knowledge_bases parameter +# Make a completion request with vector_store_ids parameter response = await litellm.acompletion( model="anthropic/claude-3-5-sonnet", messages=[{"role": "user", "content": "What is litellm?"}], - knowledge_bases=["YOUR_KNOWLEDGE_BASE_ID"] # e.g., "T37J8R4WTM" + vector_store_ids=["YOUR_KNOWLEDGE_BASE_ID"] # e.g., "T37J8R4WTM" ) print(response.choices[0].message.content) @@ -36,7 +36,7 @@ model_list: ``` -#### 2. Make a request with knowledge_bases parameter +#### 2. Make a request with vector_store_ids parameter import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; @@ -51,7 +51,7 @@ curl http://localhost:4000/v1/chat/completions \ -d '{ "model": "claude-3-5-sonnet", "messages": [{"role": "user", "content": "What is litellm?"}], - "knowledge_bases": ["YOUR_KNOWLEDGE_BASE_ID"] + "vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"] }' ``` @@ -68,11 +68,11 @@ client = OpenAI( api_key="your-litellm-api-key" ) -# Make a completion request with knowledge_bases parameter +# Make a completion request with vector_store_ids parameter response = client.chat.completions.create( model="claude-3-5-sonnet", messages=[{"role": "user", "content": "What is litellm?"}], - extra_body={"knowledge_bases": ["YOUR_KNOWLEDGE_BASE_ID"]} + extra_body={"vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"]} ) print(response.choices[0].message.content) @@ -85,7 +85,7 @@ print(response.choices[0].message.content) LiteLLM implements a `BedrockKnowledgeBaseHook` that intercepts your completion requests for handling the integration with Bedrock Knowledge Bases. -1. You make a completion request with the `knowledge_bases` parameter +1. You make a completion request with the `vector_store_ids` parameter 2. LiteLLM automatically: - Uses your last message as the query to retrieve relevant information from the Knowledge Base - Adds the retrieved context to your conversation @@ -93,7 +93,7 @@ LiteLLM implements a `BedrockKnowledgeBaseHook` that intercepts your completion ### Example Transformation -When you pass `knowledge_bases=["YOUR_KNOWLEDGE_BASE_ID"]`, your request flows through these steps: +When you pass `vector_store_ids=["YOUR_KNOWLEDGE_BASE_ID"]`, your request flows through these steps: **1. Original Request to LiteLLM:** ```json @@ -102,7 +102,7 @@ When you pass `knowledge_bases=["YOUR_KNOWLEDGE_BASE_ID"]`, your request flows t "messages": [ {"role": "user", "content": "What is litellm?"} ], - "knowledge_bases": ["YOUR_KNOWLEDGE_BASE_ID"] + "vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"] } ``` @@ -127,7 +127,7 @@ This is sent to: `https://bedrock-agent-runtime.{aws_region}.amazonaws.com/knowl } ``` -This process happens automatically whenever you include the `knowledge_bases` parameter in your request. +This process happens automatically whenever you include the `vector_store_ids` parameter in your request. ## API Reference @@ -137,4 +137,4 @@ When using the Knowledge Base integration with LiteLLM, you can include the foll | Parameter | Type | Description | |-----------|------|-------------| -| `knowledge_bases` | List[str] | List of Bedrock Knowledge Base IDs to query | +| `vector_store_ids` | List[str] | List of Bedrock Knowledge Base IDs to query | diff --git a/enterprise/proxy/enterprise_routes.py b/enterprise/proxy/enterprise_routes.py index 1d7a39b0ae5..26183874c6d 100644 --- a/enterprise/proxy/enterprise_routes.py +++ b/enterprise/proxy/enterprise_routes.py @@ -1,8 +1,10 @@ from fastapi import APIRouter from fastapi.responses import Response from .utils import _should_block_robots +from .vector_stores.endpoints import router as vector_stores_router router = APIRouter() +router.include_router(vector_stores_router) @router.get("/robots.txt") @@ -14,4 +16,4 @@ async def get_robots(): if _should_block_robots(): return Response(content="User-agent: *\nDisallow: /", media_type="text/plain") else: - return Response(status_code=404) \ No newline at end of file + return Response(status_code=404) diff --git a/enterprise/proxy/vector_stores/endpoints.py b/enterprise/proxy/vector_stores/endpoints.py new file mode 100644 index 00000000000..443f4ac3e66 --- /dev/null +++ b/enterprise/proxy/vector_stores/endpoints.py @@ -0,0 +1,180 @@ +""" +VECTOR STORE MANAGEMENT + +All /vector_store management endpoints + +/vector_store/new +/vector_store/delete +/vector_store/list +""" +import copy +from typing import List +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from fastapi import APIRouter, Depends, HTTPException +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.vector_stores.vector_store_registry import global_vector_store_manager +from litellm.types.vector_stores import ( + LiteLLM_ManagedVectorStoreListResponse, + LiteLLM_ManagedVectorStore, + VectorStoreDeleteRequest, +) + +router = APIRouter() + + +@router.post( + "/vector_store/new", + tags=["vector store management"], + dependencies=[Depends(user_api_key_auth)], +) +async def new_vector_store( + vector_store: LiteLLM_ManagedVectorStore, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Create a new vector store. + + Parameters: + - vector_store_id: str - Unique identifier for the vector store + - custom_llm_provider: str - Provider of the vector store + - vector_store_name: Optional[str] - Name of the vector store + - vector_store_description: Optional[str] - Description of the vector store + - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + # Check if vector store already exists + existing_vector_store = await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store.get("vector_store_id")} + ) + if existing_vector_store is not None: + raise HTTPException( + status_code=400, + detail=f"Vector store with ID {vector_store.get('vector_store_id')} already exists" + ) + + if vector_store.get("vector_store_metadata") is not None: + vector_store["vector_store_metadata"] = safe_dumps(vector_store.get("vector_store_metadata")) + + new_vector_store = await prisma_client.db.litellm_managedvectorstorestable.create( + data=vector_store + ) + + return { + "status": "success", + "message": f"Vector store {vector_store.get('vector_store_id')} created successfully", + "vector_store": new_vector_store, + } + except Exception as e: + verbose_proxy_logger.exception(f"Error creating vector store: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + + +@router.get( + "/vector_store/list", + tags=["vector store management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_ManagedVectorStoreListResponse, +) +async def list_vector_stores( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + page: int = 1, + page_size: int = 100, +): + """ + List all available vector stores with optional filtering and pagination. + Combines both in-memory vector stores and those stored in the database. + + Parameters: + - page: int - Page number for pagination (default: 1) + - page_size: int - Number of items per page (default: 100) + """ + from litellm.proxy.proxy_server import prisma_client + + try: + # Get in-memory vector stores + in_memory_vector_stores: List[LiteLLM_ManagedVectorStore] = copy.deepcopy(global_vector_store_manager.vector_stores) + + # Get vector stores from database + vector_stores_from_db: List[LiteLLM_ManagedVectorStore] = [] + if prisma_client is not None: + _vector_stores_from_db = await prisma_client.db.litellm_managedvectorstorestable.find_many( + order={ + "created_at": "desc" + }, + ) + for vector_store in _vector_stores_from_db: + _dict_vector_store = dict(vector_store) + _litellm_managed_vector_store = LiteLLM_ManagedVectorStore(**_dict_vector_store) + vector_stores_from_db.append(_litellm_managed_vector_store) + + # Combine in-memory and database vector stores + combined_vector_stores: List[LiteLLM_ManagedVectorStore] = in_memory_vector_stores + vector_stores_from_db + + total_count = len(combined_vector_stores) + total_pages = (total_count + page_size - 1) // page_size + + # Format response using LiteLLM_ManagedVectorStoreListResponse + response = LiteLLM_ManagedVectorStoreListResponse( + object="list", + data=combined_vector_stores, + total_count=total_count, + current_page=page, + total_pages=total_pages, + ) + + return response + except Exception as e: + verbose_proxy_logger.exception(f"Error listing vector stores: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/vector_store/delete", + tags=["vector store management"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_vector_store( + data: VectorStoreDeleteRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete a vector store. + + Parameters: + - vector_store_id: str - ID of the vector store to delete + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + # Check if vector store exists + existing_vector_store = await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": data.vector_store_id} + ) + if existing_vector_store is None: + raise HTTPException( + status_code=404, + detail=f"Vector store with ID {data.vector_store_id} not found" + ) + + # Delete vector store + await prisma_client.db.litellm_managedvectorstorestable.delete( + where={"vector_store_id": data.vector_store_id} + ) + + return {"message": f"Vector store {data.vector_store_id} deleted successfully"} + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql new file mode 100644 index 00000000000..39e7f2f3b20 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql @@ -0,0 +1,14 @@ +-- CreateTable +CREATE TABLE "LiteLLM_ManagedVectorStoresTable" ( + "vector_store_id" TEXT NOT NULL, + "custom_llm_provider" TEXT NOT NULL, + "vector_store_name" TEXT, + "vector_store_description" TEXT, + "vector_store_metadata" JSONB, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + "litellm_credential_name" TEXT, + + CONSTRAINT "LiteLLM_ManagedVectorStoresTable_pkey" PRIMARY KEY ("vector_store_id") +); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 42bb068b918..4c9856909ce 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -425,3 +425,14 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) } + +model LiteLLM_ManagedVectorStoresTable { + vector_store_id String @id + custom_llm_provider String + vector_store_name String? + vector_store_description String? + vector_store_metadata Json? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + litellm_credential_name String? +} \ No newline at end of file diff --git a/litellm/integrations/rag_hooks/bedrock_knowledgebase.py b/litellm/integrations/rag_hooks/bedrock_knowledgebase.py index 7dddf594e7a..4bd45e6b2e7 100644 --- a/litellm/integrations/rag_hooks/bedrock_knowledgebase.py +++ b/litellm/integrations/rag_hooks/bedrock_knowledgebase.py @@ -62,11 +62,11 @@ class BedrockKnowledgeBaseHook(CustomPromptManagement, BaseAWSLLM): """ Retrieves the context from the Bedrock Knowledge Base and appends it to the messages. """ - knowledge_bases = non_default_params.pop("knowledge_bases", None) - if knowledge_bases: - for knowledge_base in knowledge_bases: + vector_store_ids = non_default_params.pop("vector_store_ids", None) + if vector_store_ids: + for vector_store_id in vector_store_ids: response = await self.make_bedrock_kb_retrieve_request( - knowledge_base_id=knowledge_base, + knowledge_base_id=vector_store_id, query=self._get_kb_query_from_messages(messages), ) verbose_logger.debug(f"Bedrock Knowledge Base Response: {response}") @@ -231,7 +231,7 @@ class BedrockKnowledgeBaseHook(CustomPromptManagement, BaseAWSLLM): @staticmethod def should_use_prompt_management_hook(non_default_params: Dict) -> bool: - if non_default_params.get("knowledge_bases", None): + if non_default_params.get("vector_store_ids", None): return True return False diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 9e56d69b90b..7fa6492f6df 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,11 +1,22 @@ model_list: - - model_name: computer-user-agent + - model_name: openai/gpt-4o litellm_params: - model: azure/computer-use-preview - truncation: auto - api_key: os.environ/AZURE_RESPONSES_OPENAI_API_KEY - api_base: os.environ/AZURE_RESPONSES_OPENAI_ENDPOINT - api_version: os.environ/AZURE_RESPONSES_OPENAI_API_VERSION -general_settings: - store_prompts_in_spend_logs: true - block_robots: true + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + - model_name: claude-3-5-sonnet-with-vector-store + litellm_params: + model: anthropic/claude-3-5-sonnet-latest + api_key: os.environ/ANTHROPIC_API_KEY + vector_store_ids: ["T37J8R4WTM"] + +vector_stores: + - vector_store_name: "bedrock-litellm-website-knowledgebase" + litellm_params: + custom_llm_provider: "bedrock" + vector_store_id: "T37J8R4WTM" + vector_store_description: "Bedrock vector store for the Litellm website knowledgebase" + vector_store_metadata: + source: "https://www.litellm.com/docs" + + + diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7feb7f50929..c056147addd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1998,6 +1998,19 @@ class ProxyConfig: all_guardrails=guardrails_v2, config_file_path=config_file_path ) + ## CREDENTIALS + credential_list_dict = self.load_credential_list(config=config) + litellm.credential_list = credential_list_dict + + ## NON-LLM CONFIGS eg. MCP tools, vector stores, etc. + self._init_non_llm_configs(config=config) + + return router, router.get_model_list(), general_settings + + def _init_non_llm_configs(self, config: dict): + """ + Initialize non-LLM configs eg. MCP tools, vector stores, etc. + """ ## MCP TOOLS mcp_tools_config = config.get("mcp_tools", None) if mcp_tools_config: @@ -2011,10 +2024,17 @@ class ProxyConfig: global_mcp_server_manager.load_servers_from_config(mcp_servers_config) - ## CREDENTIALS - credential_list_dict = self.load_credential_list(config=config) - litellm.credential_list = credential_list_dict - return router, router.get_model_list(), general_settings + ## VECTOR STORES + vector_stores_config = config.get("vector_stores", None) + if vector_stores_config: + from litellm.proxy.vector_stores.vector_store_registry import ( + global_vector_store_manager, + ) + + global_vector_store_manager.load_vector_stores_from_config( + vector_stores_config + ) + pass def _load_alerting_settings(self, general_settings: dict): """ diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 42bb068b918..4c9856909ce 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -425,3 +425,14 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) } + +model LiteLLM_ManagedVectorStoresTable { + vector_store_id String @id + custom_llm_provider String + vector_store_name String? + vector_store_description String? + vector_store_metadata Json? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + litellm_credential_name String? +} \ No newline at end of file diff --git a/litellm/proxy/vector_stores/vector_store_registry.py b/litellm/proxy/vector_stores/vector_store_registry.py new file mode 100644 index 00000000000..3eb337b7724 --- /dev/null +++ b/litellm/proxy/vector_stores/vector_store_registry.py @@ -0,0 +1,79 @@ +# litellm/proxy/vector_stores/vector_store_registry.py +import json +from datetime import datetime, timezone +from typing import Any, Dict, List + +from litellm._logging import verbose_logger +from litellm.types.vector_stores import ( + LiteLLM_ManagedVectorStore, + LiteLLM_ManagedVectorStoreListResponse, + LiteLLM_VectorStoreConfig, +) + + +class VectorStoreManager: + def __init__(self): + self.vector_stores: List[LiteLLM_ManagedVectorStore] = [] + + def load_vector_stores_from_config(self, vector_stores_config: List[Dict]): + for vector_store_config in vector_stores_config: + # cast to VectorStoreConfig + litellm_vector_store_config = LiteLLM_VectorStoreConfig( + **vector_store_config + ) + vector_store_name = litellm_vector_store_config.get("vector_store_name") + vector_store_litellm_params: Dict[str, Any] = ( + litellm_vector_store_config.get("litellm_params") or {} + ) + + vector_store_id = vector_store_litellm_params.get("vector_store_id") + if vector_store_id is None: + raise ValueError( + f"vector_store_id is required for initializing vector store, got vector_store_id={vector_store_id}" + ) + custom_llm_provider = vector_store_litellm_params.get("custom_llm_provider") + if custom_llm_provider is None: + raise ValueError( + f"custom_llm_provider is required for initializing vector store, got custom_llm_provider={custom_llm_provider}" + ) + + litellm_managed_vector_store = LiteLLM_ManagedVectorStore( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider, + vector_store_name=vector_store_name, + vector_store_description=vector_store_litellm_params.get( + "vector_store_description" + ), + vector_store_metadata=vector_store_litellm_params.get( + "vector_store_metadata" + ), + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + self.vector_stores.append(litellm_managed_vector_store) + + verbose_logger.debug( + "all loaded vector stores = %s", + json.dumps(self.vector_stores, indent=4, default=str), + ) + + def list_all_vector_stores(self) -> LiteLLM_ManagedVectorStoreListResponse: + """ + List all vector stores in the required format + + Returns: + LiteLLM_ManagedVectorStoreListResponse: A standardized response with vector store data + """ + # Prepare the response + response = LiteLLM_ManagedVectorStoreListResponse( + object="list", + data=self.vector_stores, + total_count=len(self.vector_stores), + current_page=1, + total_pages=1, + ) + + return response + + +global_vector_store_manager = VectorStoreManager() diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 18e5ebd43a3..4be152c6e49 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2272,6 +2272,7 @@ class DynamicPromptManagementParamLiteral(str, Enum): CACHE_CONTROL_INJECTION_POINTS = "cache_control_injection_points" KNOWLEDGE_BASES = "knowledge_bases" + VECTOR_STORE_IDS = "vector_store_ids" @classmethod def list_all_params(cls): diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py new file mode 100644 index 00000000000..38419717750 --- /dev/null +++ b/litellm/types/vector_stores.py @@ -0,0 +1,63 @@ +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Literal, Optional, Union + +from pydantic import BaseModel +from typing_extensions import TypedDict + +from litellm.types.router import CredentialLiteLLMParams + + +class SupportedVectorStoreIntegrations(str, Enum): + """Supported vector store integrations.""" + + BEDROCK = "bedrock" + + +class LiteLLM_VectorStoreConfig(TypedDict, total=False): + """Parameters for initializing a vector store on Litellm proxy config.yaml""" + + vector_store_name: str + litellm_params: Optional[Dict[str, Any]] + + +class LiteLLM_ManagedVectorStore(TypedDict, total=False): + """LiteLLM managed vector store object - this is is the object stored in the database""" + + vector_store_id: str + custom_llm_provider: str + + vector_store_name: Optional[str] + vector_store_description: Optional[str] + vector_store_metadata: Optional[Union[Dict[str, Any], str]] + created_at: Optional[datetime] + updated_at: Optional[datetime] + + # credential fields + litellm_credential_name: Optional[str] + + +class LiteLLM_ManagedVectorStoreListResponse(TypedDict, total=False): + """Response format for listing vector stores""" + + object: Literal["list"] # Always "list" + data: List[LiteLLM_ManagedVectorStore] + total_count: Optional[int] + current_page: Optional[int] + total_pages: Optional[int] + + +class VectorStoreUpdateRequest(BaseModel): + vector_store_id: str + custom_llm_provider: Optional[str] = None + vector_store_name: Optional[str] = None + vector_store_description: Optional[str] = None + vector_store_metadata: Optional[Dict] = None + + +class VectorStoreDeleteRequest(BaseModel): + vector_store_id: str + + +class VectorStoreInfoRequest(BaseModel): + vector_store_id: str diff --git a/schema.prisma b/schema.prisma index 42bb068b918..4c9856909ce 100644 --- a/schema.prisma +++ b/schema.prisma @@ -425,3 +425,14 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) } + +model LiteLLM_ManagedVectorStoresTable { + vector_store_id String @id + custom_llm_provider String + vector_store_name String? + vector_store_description String? + vector_store_metadata Json? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + litellm_credential_name String? +} \ No newline at end of file diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 4575f602162..4ef9263d230 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -49,7 +49,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_completion(): response = await litellm.acompletion( model="anthropic/claude-3.5-sonnet", messages=[{"role": "user", "content": "what is litellm?"}], - knowledge_bases = [ + vector_store_ids = [ "T37J8R4WTM" ], client=client @@ -90,7 +90,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call(): response = await litellm.acompletion( model="anthropic/claude-3-5-haiku-latest", messages=[{"role": "user", "content": "what is litellm?"}], - knowledge_bases = [ + vector_store_ids = [ "T37J8R4WTM" ], client=async_client @@ -116,7 +116,7 @@ async def test_openai_with_knowledge_base_mock_openai(): await litellm.acompletion( model="gpt-4", messages=[{"role": "user", "content": "what is litellm?"}], - knowledge_bases = [ + vector_store_ids = [ "T37J8R4WTM" ], client=client, diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 70387046401..c3ff7c4c70f 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -34,6 +34,7 @@ import { fetchUserModels } from "@/components/create_key_button"; import { fetchTeams } from "@/components/common_components/fetch_teams"; import MCPToolsViewer from "@/components/mcp_tools"; import TagManagement from "@/components/tag_management"; +import VectorStoreManagement from "@/components/vector_store_management"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { cx } from '@/lib/cva.config'; @@ -408,6 +409,12 @@ export default function CreateKeyPage() { userRole={userRole} userID={userID} /> + ) : page == "vector-stores" ? ( + ) : page == "new_usage" ? ( = ({ const [isLoading, setIsLoading] = useState(false); const abortControllerRef = useRef(null); const [selectedTags, setSelectedTags] = useState([]); + const [selectedVectorStores, setSelectedVectorStores] = useState([]); const [messageTraceId, setMessageTraceId] = useState(null); const chatEndRef = useRef(null); @@ -328,7 +332,8 @@ const ChatUI: React.FC = ({ updateReasoningContent, updateTimingData, updateUsageData, - traceId + traceId, + selectedVectorStores.length > 0 ? selectedVectorStores : undefined ); } else if (endpointType === EndpointType.IMAGE) { // For image generation @@ -354,7 +359,8 @@ const ChatUI: React.FC = ({ updateReasoningContent, updateTimingData, updateUsageData, - traceId + traceId, + selectedVectorStores.length > 0 ? selectedVectorStores : undefined ); } } @@ -503,6 +509,27 @@ const ChatUI: React.FC = ({ accessToken={accessToken || ""} /> + +
+ + Vector Store + + Select vector store(s) to use for this LLM API call. You can set up your vector store here. + + }> + + + + +
+ Vector Store ID: {vectorStoreDetails.vector_store_id} + {vectorStoreDetails.vector_store_description || "No description"} + + {is_admin && !isEditing && ( + + )} + + + {isEditing ? ( + + + + + + + + + + + + + + + + Provider{' '} + + + + + } + name="custom_llm_provider" + rules={[{ required: true, message: "Please select a provider" }]} + > + + {Object.entries(Providers).map(([providerEnum, providerDisplayName]) => { + // Currently only showing Bedrock since it's the only supported provider + if (providerEnum === 'Bedrock') { + return ( + +
+ {`${providerEnum} { + // Create a div with provider initial as fallback + const target = e.target as HTMLImageElement; + const parent = target.parentElement; + if (parent) { + const fallbackDiv = document.createElement('div'); + fallbackDiv.className = 'w-5 h-5 rounded-full bg-gray-200 flex items-center justify-center text-xs'; + fallbackDiv.textContent = providerDisplayName.charAt(0); + parent.replaceChild(fallbackDiv, target); + } + }} + /> + {providerDisplayName} +
+ + ); + } + return null; + })} + + + + {/* Credentials */} +
+ + Either select existing credentials OR enter provider credentials below + +
+ + + + (option?.label ?? '').toLowerCase().includes(input.toLowerCase()) + } + options={[ + { value: null, label: 'None' }, + ...credentials.map((credential) => ({ + value: credential.credential_name, + label: credential.credential_name + })) + ]} + allowClear + /> + + +
+
+ OR +
+
+ + + Metadata{' '} + + + + + } + > + setMetadataString(e.target.value)} + placeholder='{"key": "value"}' + /> + + +
+ setIsEditing(false)}>Cancel + Save Changes +
+ + + ) : ( +
+ + Vector Store Details +
+
+ ID + {vectorStoreDetails.vector_store_id} +
+
+ Name + {vectorStoreDetails.vector_store_name || "-"} +
+
+ Description + {vectorStoreDetails.vector_store_description || "-"} +
+
+ Provider +
+ {(() => { + const provider = vectorStoreDetails.custom_llm_provider || "bedrock"; + const { displayName, logo } = (() => { + // Find the enum key by matching provider_map values + const enumKey = Object.keys(provider_map).find( + key => provider_map[key].toLowerCase() === provider.toLowerCase() + ); + + if (!enumKey) { + return { displayName: provider, logo: "" }; + } + + // Get the display name from Providers enum and logo from map + const displayName = Providers[enumKey as keyof typeof Providers]; + const logo = providerLogoMap[displayName]; + + return { displayName, logo }; + })(); + + return ( + <> + {logo && ( + {`${displayName} { + const target = e.target as HTMLImageElement; + const parent = target.parentElement; + if (parent) { + const fallbackDiv = document.createElement('div'); + fallbackDiv.className = 'w-5 h-5 rounded-full bg-gray-200 flex items-center justify-center text-xs'; + fallbackDiv.textContent = displayName.charAt(0); + parent.replaceChild(fallbackDiv, target); + } + }} + /> + )} + {displayName} + + ); + })()} +
+
+
+ Metadata +
+ 
{metadataString}
+
+
+
+ Created + {vectorStoreDetails.created_at ? new Date(vectorStoreDetails.created_at).toLocaleString() : "-"} +
+
+ Last Updated + {vectorStoreDetails.updated_at ? new Date(vectorStoreDetails.updated_at).toLocaleString() : "-"} +
+
+ +
+ )} + + ); +}; + +export default VectorStoreInfoView; \ No newline at end of file