forked from phoenix/litellm-mirror
119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
"""
|
|
Test Vertex AI Pass Through
|
|
|
|
1. use Credentials client side, Assert SpendLog was created
|
|
"""
|
|
|
|
import vertexai
|
|
from vertexai.preview.generative_models import GenerativeModel
|
|
import tempfile
|
|
import json
|
|
import os
|
|
import pytest
|
|
import asyncio
|
|
|
|
|
|
# Path to your service account JSON file
|
|
SERVICE_ACCOUNT_FILE = "path/to/your/service-account.json"
|
|
|
|
|
|
def load_vertex_ai_credentials():
|
|
# Define the path to the vertex_key.json file
|
|
print("loading vertex ai credentials")
|
|
filepath = os.path.dirname(os.path.abspath(__file__))
|
|
vertex_key_path = filepath + "/vertex_key.json"
|
|
|
|
# Read the existing content of the file or create an empty dictionary
|
|
try:
|
|
with open(vertex_key_path, "r") as file:
|
|
# Read the file content
|
|
print("Read vertexai file path")
|
|
content = file.read()
|
|
|
|
# If the file is empty or not valid JSON, create an empty dictionary
|
|
if not content or not content.strip():
|
|
service_account_key_data = {}
|
|
else:
|
|
# Attempt to load the existing JSON content
|
|
file.seek(0)
|
|
service_account_key_data = json.load(file)
|
|
except FileNotFoundError:
|
|
# If the file doesn't exist, create an empty dictionary
|
|
service_account_key_data = {}
|
|
|
|
# Update the service_account_key_data with environment variables
|
|
private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "")
|
|
private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "")
|
|
private_key = private_key.replace("\\n", "\n")
|
|
service_account_key_data["private_key_id"] = private_key_id
|
|
service_account_key_data["private_key"] = private_key
|
|
|
|
# Create a temporary file
|
|
with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file:
|
|
# Write the updated content to the temporary files
|
|
json.dump(service_account_key_data, temp_file, indent=2)
|
|
|
|
# Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS
|
|
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name)
|
|
|
|
|
|
async def call_spend_logs_endpoint():
|
|
"""
|
|
Call this
|
|
curl -X GET "http://0.0.0.0:4000/spend/logs" -H "Authorization: Bearer sk-1234"
|
|
"""
|
|
import datetime
|
|
import requests
|
|
|
|
todays_date = datetime.datetime.now().strftime("%Y-%m-%d")
|
|
url = f"http://0.0.0.0:4000/global/spend/logs?api_key=best-api-key-ever"
|
|
headers = {"Authorization": f"Bearer sk-1234"}
|
|
response = requests.get(url, headers=headers)
|
|
print("response from call_spend_logs_endpoint", response)
|
|
|
|
json_response = response.json()
|
|
|
|
# get spend for today
|
|
"""
|
|
json response looks like this
|
|
|
|
[{'date': '2024-08-30', 'spend': 0.00016600000000000002, 'api_key': 'best-api-key-ever'}]
|
|
"""
|
|
|
|
todays_date = datetime.datetime.now().strftime("%Y-%m-%d")
|
|
for spend_log in json_response:
|
|
if spend_log["date"] == todays_date:
|
|
return spend_log["spend"]
|
|
|
|
|
|
LITE_LLM_ENDPOINT = "http://localhost:4000"
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_basic_vertex_ai_pass_through_with_spendlog():
|
|
|
|
spend_before = await call_spend_logs_endpoint() or 0.0
|
|
# load_vertex_ai_credentials()
|
|
|
|
vertexai.init(
|
|
project="adroit-crow-413218",
|
|
location="us-central1",
|
|
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex-ai",
|
|
api_transport="rest",
|
|
)
|
|
|
|
model = GenerativeModel(model_name="gemini-1.0-pro")
|
|
response = model.generate_content("hi")
|
|
|
|
print("response", response)
|
|
|
|
await asyncio.sleep(20)
|
|
spend_after = await call_spend_logs_endpoint()
|
|
print("spend_after", spend_after)
|
|
assert (
|
|
spend_after > spend_before
|
|
), "Spend should be greater than before. spend_before: {}, spend_after: {}".format(
|
|
spend_before, spend_after
|
|
)
|
|
|
|
pass
|