From f947cec7fcdfe4a5ac2b39c8d663e04860f07f1e Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 21 Aug 2024 17:05:47 -0700
Subject: [PATCH] add test vtx embedding

---
 litellm/proxy/tests/test_vtx_embedding.py     | 21 +++++++
 litellm/proxy/tests/test_vtx_sdk_embedding.py | 59 +++++++++++++++++++
 2 files changed, 80 insertions(+)
 create mode 100644 litellm/proxy/tests/test_vtx_embedding.py
 create mode 100644 litellm/proxy/tests/test_vtx_sdk_embedding.py

diff --git a/litellm/proxy/tests/test_vtx_embedding.py b/litellm/proxy/tests/test_vtx_embedding.py
new file mode 100644
index 0000000000..4c770ae2e9
--- /dev/null
+++ b/litellm/proxy/tests/test_vtx_embedding.py
@@ -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)
diff --git a/litellm/proxy/tests/test_vtx_sdk_embedding.py b/litellm/proxy/tests/test_vtx_sdk_embedding.py
new file mode 100644
index 0000000000..a6468884f9
--- /dev/null
+++ b/litellm/proxy/tests/test_vtx_sdk_embedding.py
@@ -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}")