scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

[Feat] UI Vector Stores - Allow adding Vertex RAG Engine, OpenAI, Azure (#12752)

Browse source 
* fix _pass_through_endpoint_without_required_model

* add get_litellm_managed_vector_store_from_registry

* undo router change

* fix for using router + vector search methods

* add simple helper for _update_request_data_with_litellm_managed_vector_store_registry

* add vector_stores routes

* test_router_avector_store_search_passes_correct_args

* [Feat] UI - Allow clicking into Vector Stores (#12741)

* Add View Vector Store

* add /info for vector store

* fix updated_at

* allow easily testing the KB on litellm

* fix

* rename test

* test_init_vector_store_api_endpoints

* add get_vertex_ai_project

* fixes to vertex transformation for RAG Engine

* fix vectorStoreProviderFields

* Add Vertex Rag engine

* add oai, azure

* fix validate_environment

* fix provider name

* fix tester

* working vertex vector store
This commit is contained in:
Ishaan Jaff 2025-07-18 18:25:26 -07:00 • committed by GitHub
parent 5802a5bbe3
commit 81eb2fdd30
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 91 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/llms/base_llm/vector_store/transformation.py
View file

@ -31,6 +31,7 @@ class BaseVectorStoreConfig:
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> Tuple[str, Dict]:
pass

1
litellm/llms/bedrock/vector_stores/transformation.py
View file

@ -60,6 +60,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> Tuple[str, Dict]:
if isinstance(query, list):
query = " ".join(query)

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -2700,6 +2700,7 @@ class BaseLLMHTTPHandler:
vector_store_search_optional_params=vector_store_search_optional_params,
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
)
)
all_optional_params: Dict[str, Any] = dict(litellm_params)
@ -2796,6 +2797,7 @@ class BaseLLMHTTPHandler:
vector_store_search_optional_params=vector_store_search_optional_params,
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
)
)

2
litellm/llms/openai/vector_stores/transformation.py
View file

@ -39,6 +39,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
headers.update(
{
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
)
@ -83,6 +84,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> Tuple[str, Dict]:
url = f"{api_base}/{vector_store_id}/search"
typed_request_body = VectorStoreSearchRequest(

2
litellm/llms/pg_vector/vector_stores/transformation.py
View file

@ -81,6 +81,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> Tuple[str, Dict]:
url = f"{api_base}/{vector_store_id}/search"
_, request_body = super().transform_search_vector_store_request(
@ -89,5 +90,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
vector_store_search_optional_params=vector_store_search_optional_params,
api_base=api_base,
litellm_logging_obj=litellm_logging_obj,
litellm_params=litellm_params,
)
return url, request_body

35
litellm/llms/vertex_ai/vector_stores/transformation.py
View file

@ -42,16 +42,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
litellm_params = litellm_params or GenericLiteLLMParams()
# Get credentials and project info
vertex_credentials = (
litellm_params.vertex_credentials
or get_secret_str("VERTEXAI_CREDENTIALS")
or get_secret_str("VERTEX_AI_CREDENTIALS")
)
vertex_project = (
litellm_params.vertex_project
or get_secret_str("VERTEXAI_PROJECT")
or get_secret_str("VERTEX_AI_PROJECT")
)
vertex_credentials = self.get_vertex_ai_credentials(dict(litellm_params))
vertex_project = self.get_vertex_ai_project(dict(litellm_params))
# Get access token using the base class method
access_token, project_id = self._ensure_access_token(
@ -75,17 +67,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
"""
Get the Base endpoint for Vertex AI RAG API
"""
vertex_location = (
litellm_params.get("vertex_location")
or get_secret_str("VERTEXAI_LOCATION")
or get_secret_str("VERTEX_AI_LOCATION")
or "us-central1"
)
vertex_project = (
litellm_params.get("vertex_project")
or get_secret_str("VERTEXAI_PROJECT")
or get_secret_str("VERTEX_AI_PROJECT")
)
vertex_location = self.get_vertex_ai_location(litellm_params)
vertex_project = self.get_vertex_ai_project(litellm_params)
if api_base:
return api_base.rstrip("/")
@ -100,6 +83,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> Tuple[str, Dict[str, Any]]:
"""
Transform search request for Vertex AI RAG API
@ -111,12 +95,19 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
# Vertex AI RAG API endpoint for retrieving contexts
url = f"{api_base}:retrieveContexts"
# Use helper methods to get project and location, then construct full rag corpus path
vertex_project = self.get_vertex_ai_project(litellm_params)
vertex_location = self.get_vertex_ai_location(litellm_params)
# Construct full rag corpus path
full_rag_corpus = f"projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{vector_store_id}"
# Build the request body for Vertex AI RAG API
request_body: Dict[str, Any] = {
"vertex_rag_store": {
"rag_resources": [
{
"rag_corpus": vector_store_id
"rag_corpus": full_rag_corpus
}
]
},

1
tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py
View file

@ -20,6 +20,7 @@ def test_transform_search_request():
vector_store_search_optional_params={},
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
)
assert url.endswith("/kb123/retrieve")

2
tests/vector_store_tests/test_vertex_ai_vector_store.py
View file

@ -19,7 +19,7 @@ class TestVertexAIVectorStore(BaseVectorStoreTest):
def get_base_request_args(self):
return {
"vector_store_id": "projects/reliablekeys/locations/us-central1/ragCorpora/6917529027641081856",
"vector_store_id": "6917529027641081856",
"custom_llm_provider": "vertex_ai",
"vertex_project": "reliablekeys",
"vertex_location": "us-central1",

29
ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx
View file

@ -164,6 +164,27 @@ const VectorStoreForm: React.FC<VectorStoreFormProps> = ({
/>
)}
{/* Vertex RAG Engine Setup Instructions */}
{selectedProvider === "vertex_rag_engine" && (
<Alert
message="Vertex AI RAG Engine Setup"
description={
<div>
<p>To use Vertex AI RAG Engine:</p>
<ol style={{ marginLeft: '16px', marginTop: '8px' }}>
<li>Set up your Vertex AI RAG Engine corpus following the guide: <a href="https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/rag-overview" target="_blank" rel="noopener noreferrer">Vertex AI RAG Engine Overview</a></li>
<li>Create a corpus in your Google Cloud project</li>
<li>Note the corpus ID from the Vertex AI console</li>
<li>Enter the corpus ID in the Vector Store ID field below</li>
</ol>
</div>
}
type="info"
showIcon
style={{ marginBottom: '16px' }}
/>
)}
<Form.Item
label={
<span>
@ -176,7 +197,13 @@ const VectorStoreForm: React.FC<VectorStoreFormProps> = ({
name="vector_store_id"
rules={[{ required: true, message: "Please input the vector store ID from your api provider" }]}
>
<TextInput />
<TextInput
placeholder={
selectedProvider === "vertex_rag_engine"
? "6917529027641081856 (Get corpus ID from Vertex AI console)"
: "Enter vector store ID from your provider"
}
/>
</Form.Item>
{/* Provider-specific fields */}

42
ui/litellm-dashboard/src/components/vector_store_providers.tsx
View file

@ -1,11 +1,17 @@
export enum VectorStoreProviders {
Bedrock = "Amazon Bedrock",
PgVector = "PostgreSQL pgvector (LiteLLM Connector)"
PgVector = "PostgreSQL pgvector (LiteLLM Connector)",
VertexRagEngine = "Vertex AI RAG Engine",
OpenAI = "OpenAI",
Azure = "Azure OpenAI"
}
export const vectorStoreProviderMap: Record<string, string> = {
Bedrock: "bedrock",
PgVector: "pg_vector"
PgVector: "pg_vector",
VertexRagEngine: "vertex_ai",
OpenAI: "openai",
Azure: "azure"
};
const asset_logos_folder = '/assets/logos/';
@ -13,6 +19,9 @@ const asset_logos_folder = '/assets/logos/';
export const vectorStoreProviderLogoMap: Record<string, string> = {
[VectorStoreProviders.Bedrock]: `${asset_logos_folder}bedrock.svg`,
[VectorStoreProviders.PgVector]: `${asset_logos_folder}postgresql.svg`, // Fallback to a generic database icon if needed
[VectorStoreProviders.VertexRagEngine]: `${asset_logos_folder}google.svg`,
[VectorStoreProviders.OpenAI]: `${asset_logos_folder}openai_small.svg`,
[VectorStoreProviders.Azure]: `${asset_logos_folder}microsoft_azure.svg`,
};
// Define field types for provider-specific configurations
@ -45,6 +54,35 @@ export const vectorStoreProviderFields: Record<string, VectorStoreFieldConfig[]>
required: true,
type: "password"
}
],
vertex_rag_engine: [],
openai: [
{
name: "api_key",
label: "API Key",
tooltip: "Enter your OpenAI API key",
placeholder: "sk-...",
required: true,
type: "password"
}
],
azure: [
{
name: "api_key",
label: "API Key",
tooltip: "Enter your Azure OpenAI API key",
placeholder: "your-azure-api-key",
required: true,
type: "password"
},
{
name: "api_base",
label: "API Base",
tooltip: "Enter your Azure OpenAI endpoint (e.g., https://your-resource.openai.azure.com/)",
placeholder: "https://your-resource.openai.azure.com/",
required: true,
type: "text"
}
]
};