From bd68714f51b2575ca89691945f4da55c127c2dec Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 31 Jul 2024 15:15:18 -0700
Subject: [PATCH] fix(utils.py): map cohere timeout error

---
 litellm/tests/test_completion.py | 2 ++
 litellm/utils.py                 | 7 +++++++
 2 files changed, 9 insertions(+)

diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 6e26c28c8d..99542a9a20 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -719,6 +719,8 @@ def test_completion_cohere_command_r_plus_function_call():
             force_single_step=True,
         )
         print(second_response)
+    except litellm.Timeout:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
diff --git a/litellm/utils.py b/litellm/utils.py
index fe65e415b2..02ec7415a9 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -7328,6 +7328,13 @@ def exception_type(
                             model=model,
                             response=original_exception.response,
                         )
+                    elif original_exception.status_code == 408:
+                        exception_mapping_worked = True
+                        raise Timeout(
+                            message=f"CohereException - {original_exception.message}",
+                            llm_provider="cohere",
+                            model=model,
+                        )
                     elif original_exception.status_code == 500:
                         exception_mapping_worked = True
                         raise ServiceUnavailableError(