phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(feat) router, add ModelResponse type hints

Browse source
This commit is contained in:
ishaan-jaff 2023-12-30 10:44:13 +05:30
parent 44fc20a0ff
commit 459ba5b45e
1 changed files with 4 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

5
litellm/router.py
View file

@ -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?"}]