From deff357c923adc2327854d08a0f768029253a534 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 27 Aug 2024 14:29:18 -0700
Subject: [PATCH] fix(router.py): fix aembedding type hints

Fixes https://github.com/BerriAI/litellm/issues/5383
---
 litellm/router.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index a263d4cca4a..2d180f07637 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -1930,7 +1930,7 @@ class Router:
         input: Union[str, List],
         is_async: Optional[bool] = False,
         **kwargs,
-    ) -> Union[List[float], None]:
+    ) -> litellm.EmbeddingResponse:
         try:
             kwargs["model"] = model
             kwargs["input"] = input
@@ -2018,7 +2018,7 @@ class Router:
         input: Union[str, List],
         is_async: Optional[bool] = True,
         **kwargs,
-    ) -> Union[List[float], None]:
+    ) -> litellm.EmbeddingResponse:
         try:
             kwargs["model"] = model
             kwargs["input"] = input