update openai setTimeout

microsoft · Dec 5, 2023 · 9d2a813 · 9d2a813
1 parent a86c03f
commit 9d2a813
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 13 deletions.
diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAI.scala
@@ -4,7 +4,7 @@
 package com.microsoft.azure.synapse.ml.services.openai
 
 import com.microsoft.azure.synapse.ml.codegen.GenerationUtils
-import com.microsoft.azure.synapse.ml.services.{HasAPIVersion, HasServiceParams}
+import com.microsoft.azure.synapse.ml.services.{CognitiveServicesBase, HasAPIVersion, HasServiceParams}
 import com.microsoft.azure.synapse.ml.param.ServiceParam
 import org.apache.spark.sql.Row
 import spray.json.DefaultJsonProtocol._
@@ -244,3 +244,6 @@ trait HasOpenAITextParams extends HasOpenAISharedParams {
   }
 }
 
+abstract class OpenAIServicesBase(override val uid: String) extends CognitiveServicesBase(uid: String) {
+  setDefault(timeout -> 360.0)
+}
diff --git a/.../src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIChatCompletion.scala b/.../src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIChatCompletion.scala
@@ -20,13 +20,11 @@ import scala.language.existentials
 
 object OpenAIChatCompletion extends ComplexParamsReadable[OpenAIChatCompletion]
 
-class OpenAIChatCompletion(override val uid: String) extends CognitiveServicesBase(uid)
+class OpenAIChatCompletion(override val uid: String) extends OpenAIServicesBase(uid)
   with HasOpenAITextParams with HasCognitiveServiceInput
   with HasInternalJsonOutputParser with SynapseMLLogging {
   logClass(FeatureNames.AiServices.OpenAI)
 
-  setDefault(timeout -> 360.0)
-
   val messagesCol: Param[String] = new Param[String](
     this, "messagesCol", "The column messages to generate chat completions for," +
       " in the chat format. This column should have type Array(Struct(role: String, content: String)).")

diff --git a/...tive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAICompletion.scala b/...tive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAICompletion.scala
@@ -19,13 +19,11 @@ import scala.language.existentials
 
 object OpenAICompletion extends ComplexParamsReadable[OpenAICompletion]
 
-class OpenAICompletion(override val uid: String) extends CognitiveServicesBase(uid)
+class OpenAICompletion(override val uid: String) extends OpenAIServicesBase(uid)
   with HasOpenAITextParams with HasPromptInputs with HasCognitiveServiceInput
   with HasInternalJsonOutputParser with SynapseMLLogging {
   logClass(FeatureNames.AiServices.OpenAI)
 
-  setDefault(timeout -> 360.0)
-
   def this() = this(Identifiable.randomUID("OpenAICompletion"))
 
   def urlPath: String = ""

diff --git a/...itive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIEmbedding.scala b/...itive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIEmbedding.scala
@@ -22,12 +22,10 @@ import scala.language.existentials
 
 object OpenAIEmbedding extends ComplexParamsReadable[OpenAIEmbedding]
 
-class OpenAIEmbedding (override val uid: String) extends CognitiveServicesBase(uid)
+class OpenAIEmbedding (override val uid: String) extends OpenAIServicesBase(uid)
   with HasOpenAISharedParams with HasCognitiveServiceInput  with SynapseMLLogging {
   logClass(FeatureNames.AiServices.OpenAI)
 
-  setDefault(timeout -> 360.0)
-
   def this() = this(Identifiable.randomUID("OpenAIEmbedding"))
 
   def urlPath: String = ""

diff --git a/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPrompt.scala b/cognitive/src/main/scala/com/microsoft/azure/synapse/ml/services/openai/OpenAIPrompt.scala
@@ -26,8 +26,6 @@ class OpenAIPrompt(override val uid: String) extends Transformer
   with HasSubscriptionKey with HasAADToken with HasCustomAuthHeader
   with ComplexParamsWritable with SynapseMLLogging {
 
-  setDefault(timeout -> 360.0)
-
   logClass(FeatureNames.AiServices.OpenAI)
 
   def this() = this(Identifiable.randomUID("OpenAIPrompt"))
@@ -75,7 +73,8 @@ class OpenAIPrompt(override val uid: String) extends Transformer
     postProcessingOptions -> Map.empty,
     outputCol -> (this.uid + "_output"),
     errorCol -> (this.uid + "_error"),
-    dropPrompt -> true
+    dropPrompt -> true,
+    timeout -> 360.0
   )
 
   override def setCustomServiceName(v: String): this.type = {

diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/io/http/HTTPTransformer.scala
@@ -73,7 +73,9 @@ trait ConcurrencyParams extends Wrappable {
     case Some(v) => setConcurrentTimeout(v)
     case None => clear(concurrentTimeout)
   }
+
   setDefault(concurrency -> 1, timeout -> 60.0)
+
 }
 
 trait HasURL extends Params {