forked from phoenix/litellm-mirror
feat - allow sending tags on vertex pass through requests (#6876)
* feat - allow tagging vertex JS SDK request * add unit testing for passing headers for pass through endpoints * fix allow using vertex_ai as the primary way for pass through vertex endpoints * docs on vertex js pass tags * add e2e test for vertex pass through with spend tags * add e2e tests for streaming vertex JS with tags * fix vertex ai testing
This commit is contained in:
Ishaan Jaff
GitHub
parent
c73ce95c01
commit
f77bf49772
7 changed files with 548 additions and 77 deletions
|
|
@ -1,31 +1,22 @@
|
|||
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({
|
||||
project: 'adroit-crow-413218',
|
||||
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
|
||||
const requestOptions = {
|
||||
customHeaders: new Headers({
|
||||
"x-litellm-api-key": "sk-1234"
|
||||
})
|
||||
customHeaders: customHeaders,
|
||||
};
|
||||
|
||||
const generativeModel = vertexAI.getGenerativeModel(
|
||||
|
|
@ -33,7 +24,7 @@ const generativeModel = vertexAI.getGenerativeModel(
|
|||
requestOptions
|
||||
);
|
||||
|
||||
async function streamingResponse() {
|
||||
async function testModel() {
|
||||
try {
|
||||
const request = {
|
||||
contents: [{role: 'user', parts: [{text: 'How are you doing today tell me your name?'}]}],
|
||||
|
|
@ -49,20 +40,4 @@ async function streamingResponse() {
|
|||
}
|
||||
}
|
||||
|
||||
|
||||
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();
|
||||
testModel();
|
||||
|
|
@ -99,7 +99,7 @@ async def test_basic_vertex_ai_pass_through_with_spendlog():
|
|||
vertexai.init(
|
||||
project="adroit-crow-413218",
|
||||
location="us-central1",
|
||||
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex-ai",
|
||||
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex_ai",
|
||||
api_transport="rest",
|
||||
)
|
||||
|
||||
|
|
@ -131,7 +131,7 @@ async def test_basic_vertex_ai_pass_through_streaming_with_spendlog():
|
|||
vertexai.init(
|
||||
project="adroit-crow-413218",
|
||||
location="us-central1",
|
||||
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex-ai",
|
||||
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex_ai",
|
||||
api_transport="rest",
|
||||
)
|
||||
|
||||
|
|
@ -170,7 +170,7 @@ async def test_vertex_ai_pass_through_endpoint_context_caching():
|
|||
vertexai.init(
|
||||
project="adroit-crow-413218",
|
||||
location="us-central1",
|
||||
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex-ai",
|
||||
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex_ai",
|
||||
api_transport="rest",
|
||||
)
|
||||
|
||||
|
|
|
|||
194
tests/pass_through_tests/test_vertex_with_spend.test.js
Normal file
194
tests/pass_through_tests/test_vertex_with_spend.test.js
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
const { VertexAI, RequestOptions } = require('@google-cloud/vertexai');
|
||||
const fs = require('fs');
|
||||
const path = require('path');
|
||||
const os = require('os');
|
||||
const { writeFileSync } = require('fs');
|
||||
|
||||
|
||||
// Import fetch if the SDK uses it
|
||||
const originalFetch = global.fetch || require('node-fetch');
|
||||
|
||||
let lastCallId;
|
||||
|
||||
// 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://127.0.0.1:4000')) {
|
||||
url = url.replace('https://', 'http://');
|
||||
}
|
||||
console.log('Patched fetch sending request to:', url);
|
||||
|
||||
const response = await originalFetch(url, options);
|
||||
|
||||
// Store the call ID if it exists
|
||||
lastCallId = response.headers.get('x-litellm-call-id');
|
||||
|
||||
return response;
|
||||
};
|
||||
|
||||
function loadVertexAiCredentials() {
|
||||
console.log("loading vertex ai credentials");
|
||||
const filepath = path.dirname(__filename);
|
||||
const vertexKeyPath = path.join(filepath, "vertex_key.json");
|
||||
|
||||
// Initialize default empty service account data
|
||||
let serviceAccountKeyData = {};
|
||||
|
||||
// Try to read existing vertex_key.json
|
||||
try {
|
||||
const content = fs.readFileSync(vertexKeyPath, 'utf8');
|
||||
if (content && content.trim()) {
|
||||
serviceAccountKeyData = JSON.parse(content);
|
||||
}
|
||||
} catch (error) {
|
||||
// File doesn't exist or is invalid, continue with empty object
|
||||
}
|
||||
|
||||
// Update with environment variables
|
||||
const privateKeyId = process.env.VERTEX_AI_PRIVATE_KEY_ID || "";
|
||||
const privateKey = (process.env.VERTEX_AI_PRIVATE_KEY || "").replace(/\\n/g, "\n");
|
||||
|
||||
serviceAccountKeyData.private_key_id = privateKeyId;
|
||||
serviceAccountKeyData.private_key = privateKey;
|
||||
|
||||
// Create temporary file
|
||||
const tempFilePath = path.join(os.tmpdir(), `vertex-credentials-${Date.now()}.json`);
|
||||
writeFileSync(tempFilePath, JSON.stringify(serviceAccountKeyData, null, 2));
|
||||
|
||||
// Set environment variable
|
||||
process.env.GOOGLE_APPLICATION_CREDENTIALS = tempFilePath;
|
||||
}
|
||||
|
||||
// Run credential loading before tests
|
||||
// beforeAll(() => {
|
||||
// loadVertexAiCredentials();
|
||||
// });
|
||||
|
||||
|
||||
|
||||
describe('Vertex AI Tests', () => {
|
||||
test('should successfully generate non-streaming content with tags', async () => {
|
||||
const vertexAI = new VertexAI({
|
||||
project: 'adroit-crow-413218',
|
||||
location: 'us-central1',
|
||||
apiEndpoint: "127.0.0.1:4000/vertex_ai"
|
||||
});
|
||||
|
||||
const customHeaders = new Headers({
|
||||
"x-litellm-api-key": "sk-1234",
|
||||
"tags": "vertex-js-sdk,pass-through-endpoint"
|
||||
});
|
||||
|
||||
const requestOptions = {
|
||||
customHeaders: customHeaders
|
||||
};
|
||||
|
||||
const generativeModel = vertexAI.getGenerativeModel(
|
||||
{ model: 'gemini-1.0-pro' },
|
||||
requestOptions
|
||||
);
|
||||
|
||||
const request = {
|
||||
contents: [{role: 'user', parts: [{text: 'Say "hello test" and nothing else'}]}]
|
||||
};
|
||||
|
||||
const result = await generativeModel.generateContent(request);
|
||||
expect(result).toBeDefined();
|
||||
|
||||
// Use the captured callId
|
||||
const callId = lastCallId;
|
||||
console.log("Captured Call ID:", callId);
|
||||
|
||||
// Wait for spend to be logged
|
||||
await new Promise(resolve => setTimeout(resolve, 15000));
|
||||
|
||||
// Check spend logs
|
||||
const spendResponse = await fetch(
|
||||
`http://127.0.0.1:4000/spend/logs?request_id=${callId}`,
|
||||
{
|
||||
headers: {
|
||||
'Authorization': 'Bearer sk-1234'
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
const spendData = await spendResponse.json();
|
||||
console.log("spendData", spendData)
|
||||
expect(spendData).toBeDefined();
|
||||
expect(spendData[0].request_id).toBe(callId);
|
||||
expect(spendData[0].call_type).toBe('pass_through_endpoint');
|
||||
expect(spendData[0].request_tags).toEqual(['vertex-js-sdk', 'pass-through-endpoint']);
|
||||
expect(spendData[0].metadata).toHaveProperty('user_api_key');
|
||||
expect(spendData[0].model).toContain('gemini');
|
||||
expect(spendData[0].spend).toBeGreaterThan(0);
|
||||
}, 25000);
|
||||
|
||||
test('should successfully generate streaming content with tags', async () => {
|
||||
const vertexAI = new VertexAI({
|
||||
project: 'adroit-crow-413218',
|
||||
location: 'us-central1',
|
||||
apiEndpoint: "127.0.0.1:4000/vertex_ai"
|
||||
});
|
||||
|
||||
const customHeaders = new Headers({
|
||||
"x-litellm-api-key": "sk-1234",
|
||||
"tags": "vertex-js-sdk,pass-through-endpoint"
|
||||
});
|
||||
|
||||
const requestOptions = {
|
||||
customHeaders: customHeaders
|
||||
};
|
||||
|
||||
const generativeModel = vertexAI.getGenerativeModel(
|
||||
{ model: 'gemini-1.0-pro' },
|
||||
requestOptions
|
||||
);
|
||||
|
||||
const request = {
|
||||
contents: [{role: 'user', parts: [{text: 'Say "hello test" and nothing else'}]}]
|
||||
};
|
||||
|
||||
const streamingResult = await generativeModel.generateContentStream(request);
|
||||
expect(streamingResult).toBeDefined();
|
||||
|
||||
|
||||
// Add some assertions
|
||||
expect(streamingResult).toBeDefined();
|
||||
|
||||
for await (const item of streamingResult.stream) {
|
||||
console.log('stream chunk:', JSON.stringify(item));
|
||||
expect(item).toBeDefined();
|
||||
}
|
||||
|
||||
const aggregatedResponse = await streamingResult.response;
|
||||
console.log('aggregated response:', JSON.stringify(aggregatedResponse));
|
||||
expect(aggregatedResponse).toBeDefined();
|
||||
|
||||
// Use the captured callId
|
||||
const callId = lastCallId;
|
||||
console.log("Captured Call ID:", callId);
|
||||
|
||||
// Wait for spend to be logged
|
||||
await new Promise(resolve => setTimeout(resolve, 15000));
|
||||
|
||||
// Check spend logs
|
||||
const spendResponse = await fetch(
|
||||
`http://127.0.0.1:4000/spend/logs?request_id=${callId}`,
|
||||
{
|
||||
headers: {
|
||||
'Authorization': 'Bearer sk-1234'
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
const spendData = await spendResponse.json();
|
||||
console.log("spendData", spendData)
|
||||
expect(spendData).toBeDefined();
|
||||
expect(spendData[0].request_id).toBe(callId);
|
||||
expect(spendData[0].call_type).toBe('pass_through_endpoint');
|
||||
expect(spendData[0].request_tags).toEqual(['vertex-js-sdk', 'pass-through-endpoint']);
|
||||
expect(spendData[0].metadata).toHaveProperty('user_api_key');
|
||||
expect(spendData[0].model).toContain('gemini');
|
||||
expect(spendData[0].spend).toBeGreaterThan(0);
|
||||
}, 25000);
|
||||
});
|
||||
165
tests/pass_through_unit_tests/test_pass_through_unit_tests.py
Normal file
165
tests/pass_through_unit_tests/test_pass_through_unit_tests.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, Mock, patch, MagicMock
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import litellm
|
||||
from typing import AsyncGenerator
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy.pass_through_endpoints.types import EndpointType
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
PassThroughEndpointLogging,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.streaming_handler import (
|
||||
PassThroughStreamingHandler,
|
||||
)
|
||||
|
||||
from fastapi import Request
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_init_kwargs_for_pass_through_endpoint,
|
||||
_update_metadata_with_tags_in_header,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.types import PassthroughStandardLoggingPayload
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_request():
|
||||
# Create a mock request with headers
|
||||
class MockRequest:
|
||||
def __init__(self, headers=None):
|
||||
self.headers = headers or {}
|
||||
|
||||
return MockRequest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_user_api_key_dict():
|
||||
return UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
)
|
||||
|
||||
|
||||
def test_update_metadata_with_tags_in_header_no_tags(mock_request):
|
||||
"""
|
||||
No tags should be added to metadata if they do not exist in headers
|
||||
"""
|
||||
# Test when no tags are present in headers
|
||||
request = mock_request(headers={})
|
||||
metadata = {"existing": "value"}
|
||||
|
||||
result = _update_metadata_with_tags_in_header(request=request, metadata=metadata)
|
||||
|
||||
assert result == {"existing": "value"}
|
||||
assert "tags" not in result
|
||||
|
||||
|
||||
def test_update_metadata_with_tags_in_header_with_tags(mock_request):
|
||||
"""
|
||||
Tags should be added to metadata if they exist in headers
|
||||
"""
|
||||
# Test when tags are present in headers
|
||||
request = mock_request(headers={"tags": "tag1,tag2,tag3"})
|
||||
metadata = {"existing": "value"}
|
||||
|
||||
result = _update_metadata_with_tags_in_header(request=request, metadata=metadata)
|
||||
|
||||
assert result == {"existing": "value", "tags": ["tag1", "tag2", "tag3"]}
|
||||
|
||||
|
||||
def test_init_kwargs_for_pass_through_endpoint_basic(
|
||||
mock_request, mock_user_api_key_dict
|
||||
):
|
||||
"""
|
||||
Basic test for init_kwargs_for_pass_through_endpoint
|
||||
|
||||
- metadata should contain user_api_key, user_api_key_user_id, user_api_key_team_id, user_api_key_end_user_id from `mock_user_api_key_dict`
|
||||
"""
|
||||
request = mock_request()
|
||||
passthrough_payload = PassthroughStandardLoggingPayload(
|
||||
url="https://test.com",
|
||||
request_body={},
|
||||
)
|
||||
|
||||
result = _init_kwargs_for_pass_through_endpoint(
|
||||
request=request,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
passthrough_logging_payload=passthrough_payload,
|
||||
litellm_call_id="test-call-id",
|
||||
)
|
||||
|
||||
assert result["call_type"] == "pass_through_endpoint"
|
||||
assert result["litellm_call_id"] == "test-call-id"
|
||||
assert result["passthrough_logging_payload"] == passthrough_payload
|
||||
|
||||
# Check metadata
|
||||
expected_metadata = {
|
||||
"user_api_key": "test-key",
|
||||
"user_api_key_user_id": "test-user",
|
||||
"user_api_key_team_id": "test-team",
|
||||
"user_api_key_end_user_id": "test-user",
|
||||
}
|
||||
assert result["litellm_params"]["metadata"] == expected_metadata
|
||||
|
||||
|
||||
def test_init_kwargs_with_litellm_metadata(mock_request, mock_user_api_key_dict):
|
||||
"""
|
||||
Expected behavior: litellm_metadata should be merged with default metadata
|
||||
|
||||
see usage example here: https://docs.litellm.ai/docs/pass_through/anthropic_completion#send-litellm_metadata-tags
|
||||
"""
|
||||
request = mock_request()
|
||||
parsed_body = {
|
||||
"litellm_metadata": {"custom_field": "custom_value", "tags": ["tag1", "tag2"]}
|
||||
}
|
||||
passthrough_payload = PassthroughStandardLoggingPayload(
|
||||
url="https://test.com",
|
||||
request_body={},
|
||||
)
|
||||
|
||||
result = _init_kwargs_for_pass_through_endpoint(
|
||||
request=request,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
passthrough_logging_payload=passthrough_payload,
|
||||
_parsed_body=parsed_body,
|
||||
litellm_call_id="test-call-id",
|
||||
)
|
||||
|
||||
# Check that litellm_metadata was merged with default metadata
|
||||
metadata = result["litellm_params"]["metadata"]
|
||||
print("metadata", metadata)
|
||||
assert metadata["custom_field"] == "custom_value"
|
||||
assert metadata["tags"] == ["tag1", "tag2"]
|
||||
assert metadata["user_api_key"] == "test-key"
|
||||
|
||||
|
||||
def test_init_kwargs_with_tags_in_header(mock_request, mock_user_api_key_dict):
|
||||
"""
|
||||
Tags should be added to metadata if they exist in headers
|
||||
"""
|
||||
request = mock_request(headers={"tags": "tag1,tag2"})
|
||||
passthrough_payload = PassthroughStandardLoggingPayload(
|
||||
url="https://test.com",
|
||||
request_body={},
|
||||
)
|
||||
|
||||
result = _init_kwargs_for_pass_through_endpoint(
|
||||
request=request,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
passthrough_logging_payload=passthrough_payload,
|
||||
litellm_call_id="test-call-id",
|
||||
)
|
||||
|
||||
# Check that tags were added to metadata
|
||||
metadata = result["litellm_params"]["metadata"]
|
||||
print("metadata", metadata)
|
||||
assert metadata["tags"] == ["tag1", "tag2"]
|
||||
Loading…
Add table
Add a link
Reference in a new issue