mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-27 11:43:54 +00:00
add test vtx embedding
This commit is contained in:
parent
e9537c6560
commit
f947cec7fc
2 changed files with 80 additions and 0 deletions
21
litellm/proxy/tests/test_vtx_embedding.py
Normal file
21
litellm/proxy/tests/test_vtx_embedding.py
Normal file
|
@ -0,0 +1,21 @@
|
||||||
|
import openai
|
||||||
|
|
||||||
|
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
|
||||||
|
|
||||||
|
# # request sent to model set on litellm proxy, `litellm --model`
|
||||||
|
response = client.embeddings.create(
|
||||||
|
model="multimodalembedding@001",
|
||||||
|
input=[],
|
||||||
|
extra_body={
|
||||||
|
"instances": [
|
||||||
|
{
|
||||||
|
"image": {
|
||||||
|
"gcsUri": "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"
|
||||||
|
},
|
||||||
|
"text": "this is a unicorn",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
print(response)
|
59
litellm/proxy/tests/test_vtx_sdk_embedding.py
Normal file
59
litellm/proxy/tests/test_vtx_sdk_embedding.py
Normal file
|
@ -0,0 +1,59 @@
|
||||||
|
import vertexai
|
||||||
|
from google.auth.credentials import Credentials
|
||||||
|
from vertexai.vision_models import (
|
||||||
|
Image,
|
||||||
|
MultiModalEmbeddingModel,
|
||||||
|
Video,
|
||||||
|
VideoSegmentConfig,
|
||||||
|
)
|
||||||
|
|
||||||
|
LITELLM_PROXY_API_KEY = "sk-1234"
|
||||||
|
LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai"
|
||||||
|
|
||||||
|
import datetime
|
||||||
|
|
||||||
|
|
||||||
|
class CredentialsWrapper(Credentials):
|
||||||
|
def __init__(self, token=None):
|
||||||
|
super().__init__()
|
||||||
|
self.token = token
|
||||||
|
self.expiry = None # or set to a future date if needed
|
||||||
|
|
||||||
|
def refresh(self, request):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def apply(self, headers, token=None):
|
||||||
|
headers["Authorization"] = f"Bearer {self.token}"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def expired(self):
|
||||||
|
return False # Always consider the token as non-expired
|
||||||
|
|
||||||
|
@property
|
||||||
|
def valid(self):
|
||||||
|
return True # Always consider the credentials as valid
|
||||||
|
|
||||||
|
|
||||||
|
credentials = CredentialsWrapper(token=LITELLM_PROXY_API_KEY)
|
||||||
|
|
||||||
|
vertexai.init(
|
||||||
|
project="adroit-crow-413218",
|
||||||
|
location="us-central1",
|
||||||
|
api_endpoint=LITELLM_PROXY_BASE,
|
||||||
|
credentials=credentials,
|
||||||
|
api_transport="rest",
|
||||||
|
request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")],
|
||||||
|
)
|
||||||
|
|
||||||
|
model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding")
|
||||||
|
image = Image.load_from_file(
|
||||||
|
"gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"
|
||||||
|
)
|
||||||
|
|
||||||
|
embeddings = model.get_embeddings(
|
||||||
|
image=image,
|
||||||
|
contextual_text="Colosseum",
|
||||||
|
dimension=1408,
|
||||||
|
)
|
||||||
|
print(f"Image Embedding: {embeddings.image_embedding}")
|
||||||
|
print(f"Text Embedding: {embeddings.text_embedding}")
|
Loading…
Add table
Add a link
Reference in a new issue