fix max token decrement and streaming logic

Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>

Author: RishabhSaini

pkg/sidecar/proxy/connector_nixlv2.go

@@ -86,12 +86,8 @@ func (s *Server) runNIXLProtocolV2(w http.ResponseWriter, r *http.Request, prefi
 	maxCompletionTokensValue, maxCompletionTokensOk := completionRequest[requestFieldMaxCompletionTokens]
 
 	// Determine if client wants streaming
-	clientWantsStreaming := false
-	if streamOk {
-		if streamBool, ok := streamValue.(bool); ok {
-			clientWantsStreaming = streamBool
-		}
-	}
+	streamBool, streamBoolOk := streamValue.(bool)
+	clientWantsStreaming := streamOk && streamBoolOk && streamBool
 
 	completionRequest[requestFieldKVTransferParams] = map[string]any{
 		requestFieldDoRemoteDecode:  true,
@@ -260,18 +256,26 @@ func (s *Server) runNIXLProtocolV2(w http.ResponseWriter, r *http.Request, prefi
 	}
 	delete(completionRequest, requestFieldMaxTokens)
 	if maxTokensOk {
-		// Decrement by 1 since prefill already generated 1 token
-		if val, ok := maxTokensValue.(float64); ok && val > 0 {
-			completionRequest[requestFieldMaxTokens] = val - 1
+		if clientWantsStreaming {
+			// Decrement by 1 since we already sent 1 token to streaming client
+			if val, ok := maxTokensValue.(float64); ok && val > 0 {
+				completionRequest[requestFieldMaxTokens] = val - 1
+			} else {
+				completionRequest[requestFieldMaxTokens] = maxTokensValue
+			}
 		} else {
 			completionRequest[requestFieldMaxTokens] = maxTokensValue
 		}
 	}
 	delete(completionRequest, requestFieldMaxCompletionTokens)
 	if maxCompletionTokensOk {
-		// Decrement by 1 since prefill already generated 1 token
-		if val, ok := maxCompletionTokensValue.(float64); ok && val > 0 {
-			completionRequest[requestFieldMaxCompletionTokens] = val - 1
+		if clientWantsStreaming {
+			// Decrement by 1 since we already sent 1 token to streaming client
+			if val, ok := maxCompletionTokensValue.(float64); ok && val > 0 {
+				completionRequest[requestFieldMaxCompletionTokens] = val - 1
+			} else {
+				completionRequest[requestFieldMaxCompletionTokens] = maxCompletionTokensValue
+			}
 		} else {
 			completionRequest[requestFieldMaxCompletionTokens] = maxCompletionTokensValue
 		}