forked from phoenix/litellm-mirror
feat - allow tagging vertex JS SDK request
This commit is contained in:
parent
772b2f9cd2
commit
19ad0cc17f
2 changed files with 24 additions and 34 deletions
|
@ -388,6 +388,7 @@ async def pass_through_request( # noqa: PLR0915
|
||||||
_parsed_body=_parsed_body,
|
_parsed_body=_parsed_body,
|
||||||
passthrough_logging_payload=passthrough_logging_payload,
|
passthrough_logging_payload=passthrough_logging_payload,
|
||||||
litellm_call_id=litellm_call_id,
|
litellm_call_id=litellm_call_id,
|
||||||
|
request=request,
|
||||||
)
|
)
|
||||||
# done for supporting 'parallel_request_limiter.py' with pass-through endpoints
|
# done for supporting 'parallel_request_limiter.py' with pass-through endpoints
|
||||||
logging_obj.update_environment_variables(
|
logging_obj.update_environment_variables(
|
||||||
|
@ -567,6 +568,7 @@ async def pass_through_request( # noqa: PLR0915
|
||||||
|
|
||||||
|
|
||||||
def _init_kwargs_for_pass_through_endpoint(
|
def _init_kwargs_for_pass_through_endpoint(
|
||||||
|
request: Request,
|
||||||
user_api_key_dict: UserAPIKeyAuth,
|
user_api_key_dict: UserAPIKeyAuth,
|
||||||
passthrough_logging_payload: PassthroughStandardLoggingPayload,
|
passthrough_logging_payload: PassthroughStandardLoggingPayload,
|
||||||
_parsed_body: Optional[dict] = None,
|
_parsed_body: Optional[dict] = None,
|
||||||
|
@ -582,6 +584,12 @@ def _init_kwargs_for_pass_through_endpoint(
|
||||||
}
|
}
|
||||||
if _litellm_metadata:
|
if _litellm_metadata:
|
||||||
_metadata.update(_litellm_metadata)
|
_metadata.update(_litellm_metadata)
|
||||||
|
|
||||||
|
_metadata = _update_metadata_with_tags_in_header(
|
||||||
|
request=request,
|
||||||
|
metadata=_metadata,
|
||||||
|
)
|
||||||
|
|
||||||
kwargs = {
|
kwargs = {
|
||||||
"litellm_params": {
|
"litellm_params": {
|
||||||
"metadata": _metadata,
|
"metadata": _metadata,
|
||||||
|
@ -593,6 +601,13 @@ def _init_kwargs_for_pass_through_endpoint(
|
||||||
return kwargs
|
return kwargs
|
||||||
|
|
||||||
|
|
||||||
|
def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> dict:
|
||||||
|
_tags = request.headers.get("tags")
|
||||||
|
if _tags:
|
||||||
|
metadata["tags"] = _tags.split(",")
|
||||||
|
return metadata
|
||||||
|
|
||||||
|
|
||||||
def create_pass_through_route(
|
def create_pass_through_route(
|
||||||
endpoint,
|
endpoint,
|
||||||
target: str,
|
target: str,
|
||||||
|
|
|
@ -1,31 +1,22 @@
|
||||||
const { VertexAI, RequestOptions } = require('@google-cloud/vertexai');
|
const { VertexAI, RequestOptions } = require('@google-cloud/vertexai');
|
||||||
|
|
||||||
|
|
||||||
// Import fetch if the SDK uses it
|
|
||||||
const originalFetch = global.fetch || require('node-fetch');
|
|
||||||
|
|
||||||
// Monkey-patch the fetch used internally
|
|
||||||
global.fetch = async function patchedFetch(url, options) {
|
|
||||||
// Modify the URL to use HTTP instead of HTTPS
|
|
||||||
if (url.startsWith('https://localhost:4000')) {
|
|
||||||
url = url.replace('https://', 'http://');
|
|
||||||
}
|
|
||||||
console.log('Patched fetch sending request to:', url);
|
|
||||||
return originalFetch(url, options);
|
|
||||||
};
|
|
||||||
|
|
||||||
const vertexAI = new VertexAI({
|
const vertexAI = new VertexAI({
|
||||||
project: 'adroit-crow-413218',
|
project: 'adroit-crow-413218',
|
||||||
location: 'us-central1',
|
location: 'us-central1',
|
||||||
apiEndpoint: "localhost:4000/vertex-ai"
|
apiEndpoint: "127.0.0.1:4000/vertex-ai"
|
||||||
});
|
});
|
||||||
|
|
||||||
|
// Create customHeaders using Headers
|
||||||
|
const customHeaders = new Headers({
|
||||||
|
"X-Litellm-Api-Key": "sk-1234",
|
||||||
|
tags: "vertexjs,test-2"
|
||||||
|
});
|
||||||
|
|
||||||
// Use customHeaders in RequestOptions
|
// Use customHeaders in RequestOptions
|
||||||
const requestOptions = {
|
const requestOptions = {
|
||||||
customHeaders: new Headers({
|
customHeaders: customHeaders,
|
||||||
"x-litellm-api-key": "sk-1234"
|
|
||||||
})
|
|
||||||
};
|
};
|
||||||
|
|
||||||
const generativeModel = vertexAI.getGenerativeModel(
|
const generativeModel = vertexAI.getGenerativeModel(
|
||||||
|
@ -33,7 +24,7 @@ const generativeModel = vertexAI.getGenerativeModel(
|
||||||
requestOptions
|
requestOptions
|
||||||
);
|
);
|
||||||
|
|
||||||
async function streamingResponse() {
|
async function testModel() {
|
||||||
try {
|
try {
|
||||||
const request = {
|
const request = {
|
||||||
contents: [{role: 'user', parts: [{text: 'How are you doing today tell me your name?'}]}],
|
contents: [{role: 'user', parts: [{text: 'How are you doing today tell me your name?'}]}],
|
||||||
|
@ -49,20 +40,4 @@ async function streamingResponse() {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
testModel();
|
||||||
async function nonStreamingResponse() {
|
|
||||||
try {
|
|
||||||
const request = {
|
|
||||||
contents: [{role: 'user', parts: [{text: 'How are you doing today tell me your name?'}]}],
|
|
||||||
};
|
|
||||||
const response = await generativeModel.generateContent(request);
|
|
||||||
console.log('non streaming response: ', JSON.stringify(response));
|
|
||||||
} catch (error) {
|
|
||||||
console.error('Error:', error);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
streamingResponse();
|
|
||||||
nonStreamingResponse();
|
|
Loading…
Add table
Add a link
Reference in a new issue