From 268702722504ec2f4bf8f72f7bb15cb6da6843b3 Mon Sep 17 00:00:00 2001
From: Minwoo Lee <11580164+minwhoo@users.noreply.github.com>
Date: Thu, 13 Feb 2025 15:40:56 +0900
Subject: [PATCH] Apply streaming-related transformations only for generate
 config

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

diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py
index 9b100ff1f8..b09f7b0444 100644
--- a/litellm/llms/triton/completion/transformation.py
+++ b/litellm/llms/triton/completion/transformation.py
@@ -52,17 +52,6 @@ class TritonConfig(BaseConfig):
     ) -> Dict:
         return {"Content-Type": "application/json"}
 
-    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 get_supported_openai_params(self, model: str) -> List:
         return ["max_tokens", "max_completion_tokens"]
 
@@ -178,6 +167,17 @@ 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,
@@ -227,7 +227,7 @@ class TritonGenerateConfig(TritonConfig):
         return model_response
 
 
-class TritonInferConfig(TritonGenerateConfig):
+class TritonInferConfig(TritonConfig):
     """
     Transformations for triton /infer endpoint (his is an infer model with a custom model on triton)
     """