open-git/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2025-12-06 11:33:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add embedding pcs support

Browse Source
This commit is contained in:
Sameer Kankute
2025-12-06 00:20:30 +05:30
parent 43914796d6
commit 64c001255d
2 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py
View File

@@ -72,6 +72,9 @@ class VertexEmbedding(VertexBase):
project_id=vertex_project,
custom_llm_provider=custom_llm_provider,
)
# Extract use_psc_endpoint_format from optional_params
use_psc_endpoint_format = optional_params.get("use_psc_endpoint_format", False)
auth_header, api_base = self._get_token_and_url(
model=model,
gemini_api_key=gemini_api_key,
@@ -84,6 +87,7 @@ class VertexEmbedding(VertexBase):
api_base=api_base,
should_use_v1beta1_features=should_use_v1beta1_features,
mode="embedding",
use_psc_endpoint_format=use_psc_endpoint_format,
)
headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
vertex_request: VertexEmbeddingRequest = (
@@ -164,6 +168,9 @@ class VertexEmbedding(VertexBase):
project_id=vertex_project,
custom_llm_provider=custom_llm_provider,
)
# Extract use_psc_endpoint_format from optional_params
use_psc_endpoint_format = optional_params.get("use_psc_endpoint_format", False)
auth_header, api_base = self._get_token_and_url(
model=model,
gemini_api_key=gemini_api_key,
@@ -176,6 +183,7 @@ class VertexEmbedding(VertexBase):
api_base=api_base,
should_use_v1beta1_features=should_use_v1beta1_features,
mode="embedding",
use_psc_endpoint_format=use_psc_endpoint_format,
)
headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
vertex_request: VertexEmbeddingRequest = (

3
tests/test_litellm/llms/vertex_ai/test_bge_embedding.py
View File

@@ -214,7 +214,8 @@ def test_vertex_ai_bge_psc_endpoint_url_construction():
api_base="http://10.128.16.2",
vertex_project="gen-lang-client-0682925754",
vertex_location="us-central1",
client=client
client=client,
use_psc_endpoint_format=True # Enable PSC endpoint format for this test
)
mock_post.assert_called_once()