scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

fix(router): satisfy lint gates for sync text_completion fix (ruff-strict, router coverage, types)

Browse source
This commit is contained in:
Lalith Vardhan 2026-10-23 10:16:11 -07:00
parent fed20b9db9
commit 11835a2a9f
2 changed files with 48 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
litellm/router.py
View file

@ -438,6 +438,7 @@ from litellm.utils import (
ModelResponse,
Rules,
TextCompletionResponse,
TextCompletionStreamWrapper,
function_setup,
get_llm_provider,
get_non_default_completion_params,
@ -5724,23 +5725,25 @@ class Router:
is_async: bool | None = False,
**kwargs,
):
try:
kwargs["model"] = model
kwargs["prompt"] = prompt
kwargs["original_function"] = self._text_completion
self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs)
return self.function_with_fallbacks(**kwargs)
except Exception as e:
raise e
kwargs["model"] = model
kwargs["prompt"] = prompt
kwargs["original_function"] = self._text_completion
self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs)
return self.function_with_fallbacks(**kwargs)
def _text_completion(self, model: str, prompt: str, **kwargs):
model_name = None
def _text_completion(
self,
model: str,
prompt: str,
**kwargs: object, # kwargs-ok: forwarded verbatim to litellm.text_completion, like _completion
) -> TextCompletionResponse | TextCompletionStreamWrapper | ModelResponse | CustomStreamWrapper:
model_name: str | None = None
try:
verbose_router_logger.debug("Inside _text_completion()- model: %s; kwargs: %s", model, kwargs)
deployment: Final = self.get_available_deployment(
model=model,
messages=[{"role": "user", "content": prompt}],
specific_deployment=kwargs.pop("specific_deployment", None),
specific_deployment=kwargs.pop("specific_deployment", None) is True,
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
@ -5756,15 +5759,15 @@ class Router:
response: Final = litellm.text_completion(
**{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs}
)
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.text_completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
return response
except Exception as e:
verbose_router_logger.info("litellm.text_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e)
if model_name is not None:
self.fail_calls[model_name] += 1
self._stamp_retry_skip_deployment_id(e, kwargs)
raise e
raise
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.text_completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
return response
async def atext_completion(
self,

30
tests/unit/test_router/test_router_endpoints.py
View file

@ -1305,3 +1305,33 @@ def test_sync_text_completion_falls_back_like_async():
response = router.text_completion(model="primary", prompt="hi")
assert response.choices[0].text == "from backup", response
def test_sync_text_completion_deployment_call_updates_call_counters():
router = Router(
model_list=[
{
"model_name": "ok-group",
"litellm_params": {"model": "gpt-3.5-turbo-instruct", "api_key": "fake-key", "mock_response": "ok"},
},
{
"model_name": "failing-group",
"litellm_params": {
"model": "openai/davinci-002",
"api_key": "fake-key",
"mock_response": "litellm.RateLimitError",
},
},
],
num_retries=0,
)
response = router._text_completion(model="ok-group", prompt="hi")
with pytest.raises(litellm.RateLimitError):
router._text_completion(model="failing-group", prompt="hi")
assert response.choices[0].text == "ok", response
assert router.total_calls["gpt-3.5-turbo-instruct"] == 1, router.total_calls
assert router.success_calls["gpt-3.5-turbo-instruct"] == 1, router.success_calls
assert router.total_calls["openai/davinci-002"] == 1, router.total_calls
assert router.fail_calls["openai/davinci-002"] == 1, router.fail_calls