From c62be184c2e9228ad321384f1c385be1ff4f882b Mon Sep 17 00:00:00 2001
From: Minwoo Lee <11580164+minwhoo@users.noreply.github.com>
Date: Thu, 13 Feb 2025 16:41:50 +0900
Subject: [PATCH] Fix get_complete_url

---
 .../llms/triton/completion/transformation.py  | 23 ++++++++++---------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py
index b09f7b04443..0a65e216dfe 100644
--- a/litellm/llms/triton/completion/transformation.py
+++ b/litellm/llms/triton/completion/transformation.py
@@ -67,6 +67,18 @@ class TritonConfig(BaseConfig):
                 optional_params[param] = value
         return optional_params
 
+    def get_complete_url(
+        self,
+        api_base: str,
+        model: str,
+        optional_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        llm_type = self._get_triton_llm_type(api_base)
+        if llm_type == "generate" and stream:
+            return api_base + "_stream"
+        return api_base
+
     def transform_response(
         self,
         model: str,
@@ -167,17 +179,6 @@ class TritonGenerateConfig(TritonConfig):
     Transformations for triton /generate endpoint (This is a trtllm model)
     """
 
-    def get_complete_url(
-        self,
-        api_base: str,
-        model: str,
-        optional_params: dict,
-        stream: Optional[bool] = None,
-    ) -> str:
-        if stream:
-            return api_base + "_stream"
-        return api_base
-
     def transform_request(
         self,
         model: str,