From 28e4706bfdcb41b13589b09b1a5d4bd955599679 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 8 Apr 2024 12:02:40 -0700
Subject: [PATCH] test - re-order embedding responses

---
 litellm/proxy/tests/test_openai_embedding.py | 126 +++++++++++++++++++
 1 file changed, 126 insertions(+)
 create mode 100644 litellm/proxy/tests/test_openai_embedding.py

diff --git a/litellm/proxy/tests/test_openai_embedding.py b/litellm/proxy/tests/test_openai_embedding.py
new file mode 100644
index 0000000000..3763f4edd7
--- /dev/null
+++ b/litellm/proxy/tests/test_openai_embedding.py
@@ -0,0 +1,126 @@
+import openai
+import asyncio
+
+
+async def async_request(client, model, input_data):
+    response = await client.embeddings.create(model=model, input=input_data)
+    response = response.dict()
+    data_list = response["data"]
+    for i, embedding in enumerate(data_list):
+        embedding["embedding"] = []
+        current_index = embedding["index"]
+        assert i == current_index
+    return response
+
+
+async def main():
+    client = openai.AsyncOpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
+    models = [
+        "text-embedding-ada-002",
+        "text-embedding-ada-002",
+        "text-embedding-ada-002",
+    ]
+    inputs = [
+        [
+            "5",
+            "6",
+            "7",
+            "8",
+            "9",
+            "10",
+            "11",
+            "12",
+            "13",
+            "14",
+            "15",
+            "16",
+            "17",
+            "18",
+            "19",
+            "20",
+        ],
+        ["1", "2", "3", "4", "5", "6"],
+        [
+            "1",
+            "2",
+            "3",
+            "4",
+            "5",
+            "6",
+            "7",
+            "8",
+            "9",
+            "10",
+            "11",
+            "12",
+            "13",
+            "14",
+            "15",
+            "16",
+            "17",
+            "18",
+            "19",
+            "20",
+        ],
+        [
+            "1",
+            "2",
+            "3",
+            "4",
+            "5",
+            "6",
+            "7",
+            "8",
+            "9",
+            "10",
+            "11",
+            "12",
+            "13",
+            "14",
+            "15",
+            "16",
+            "17",
+            "18",
+            "19",
+            "20",
+        ],
+        [
+            "1",
+            "2",
+            "3",
+            "4",
+            "5",
+            "6",
+            "7",
+            "8",
+            "9",
+            "10",
+            "11",
+            "12",
+            "13",
+            "14",
+            "15",
+            "16",
+            "17",
+            "18",
+            "19",
+            "20",
+        ],
+        ["1", "2", "3"],
+    ]
+
+    tasks = []
+    for model, input_data in zip(models, inputs):
+        task = async_request(client, model, input_data)
+        tasks.append(task)
+
+    responses = await asyncio.gather(*tasks)
+    print(responses)
+    for response in responses:
+        data_list = response["data"]
+        for embedding in data_list:
+            embedding["embedding"] = []
+        print(response)
+
+
+asyncio.run(main())