Merge 50828b36a4 into dddbaa803d

litellm/router.py

@@ -441,6 +441,7 @@ from litellm.utils import (
     ProviderConfigManager,
     Rules,
     TextCompletionResponse,
+    TextCompletionStreamWrapper,
     function_setup,
     get_llm_provider,
     get_non_default_completion_params,
@@ -5841,32 +5842,49 @@ class Router:
         is_async: bool | None = False,
         **kwargs,
     ):
-        messages: Final = [{"role": "user", "content": prompt}]
-        try:
-            kwargs["model"] = model
-            kwargs["prompt"] = prompt
-            kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries)
-            kwargs.setdefault("metadata", {}).update({"model_group": model})
+        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)
 
-            # pick the one that is available (lowest TPM/RPM)
+    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  # rebind-ok: set once the deployment is picked; read in except for fail_calls
+        try:
+            verbose_router_logger.debug("Inside _text_completion()- model: %s; kwargs: %s", model, kwargs)
             deployment: Final = self.get_available_deployment(
                 model=model,
-                messages=messages,
-                specific_deployment=kwargs.pop("specific_deployment", None),
+                messages=[{"role": "user", "content": prompt}],
+                specific_deployment=kwargs.pop("specific_deployment", None) is True,
                 request_kwargs=kwargs,
             )
+            self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
 
             data: Final = deployment["litellm_params"].copy()
-            for k, v in self.default_litellm_params.items():
-                if k not in kwargs:  # prioritize model-specific params > default router params
-                    kwargs[k] = v
-                elif k == "metadata":
-                    kwargs[k].update(v)
+            model_name = data["model"]  # rebind-ok: assigned inside try so except can attribute the failure
+            self.total_calls[model_name] += 1
 
-            # call via litellm.completion()
-            return litellm.text_completion(**{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs})
+            ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit)
+            self.routing_strategy_pre_call_checks(deployment=deployment)
+
+            # kwargs no longer carries the model group name, so the deployment's model is used
+            response: Final = litellm.text_completion(
+                **{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs}
+            )
         except Exception as e:
-            raise 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
+        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,

tests/unit/test_router/test_router_endpoints.py

@@ -1259,3 +1259,79 @@ def test_router_model_group_encrypted_content_affinity_callback_registration():
         assert encrypted_content_callbacks[0].router is router
     finally:
         router.discard()
+
+
+def test_sync_text_completion_calls_the_deployment_model_not_the_group_alias():
+    router = Router(
+        model_list=[
+            {
+                "model_name": "my-instruct-alias",
+                "litellm_params": {"model": "gpt-3.5-turbo-instruct", "api_key": "fake-key"},
+            }
+        ]
+    )
+
+    with patch.object(litellm, "text_completion", wraps=litellm.text_completion) as text_completion_spy:
+        response = router.text_completion(model="my-instruct-alias", prompt="hi", mock_response="ok")
+
+    assert response.choices[0].text == "ok", response
+    assert text_completion_spy.call_args.kwargs["model"] == "gpt-3.5-turbo-instruct", text_completion_spy.call_args
+
+
+def test_sync_text_completion_falls_back_like_async():
+    router = Router(
+        model_list=[
+            {
+                "model_name": "primary",
+                "litellm_params": {
+                    "model": "gpt-3.5-turbo-instruct",
+                    "api_key": "fake-key",
+                    "mock_response": "litellm.RateLimitError",
+                },
+            },
+            {
+                "model_name": "backup",
+                "litellm_params": {
+                    "model": "gpt-3.5-turbo-instruct",
+                    "api_key": "fake-key",
+                    "mock_response": "from backup",
+                },
+            },
+        ],
+        fallbacks=[{"primary": ["backup"]}],
+        num_retries=0,
+    )
+
+    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