From 459ba5b45e71f62b4591a34734dbba97f92793d3 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 30 Dec 2023 10:44:13 +0530 Subject: [PATCH] (feat) router, add ModelResponse type hints --- litellm/router.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index ecdc8f211..02eb91019 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -23,6 +23,7 @@ from litellm.llms.custom_httpx.azure_dall_e_2 import ( CustomHTTPTransport, AsyncCustomHTTPTransport, ) +from litellm.utils import ModelResponse, CustomStreamWrapper import copy @@ -216,7 +217,9 @@ class Router: ### COMPLETION, EMBEDDING, IMG GENERATION FUNCTIONS - def completion(self, model: str, messages: List[Dict[str, str]], **kwargs): + def completion( + self, model: str, messages: List[Dict[str, str]], **kwargs + ) -> Union[ModelResponse, CustomStreamWrapper]: """ Example usage: response = router.completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}]