diff --git a/process.go b/process.go index 65767e7..166d948 100644 --- a/process.go +++ b/process.go @@ -32,13 +32,7 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s return err } - // set timeout, default is 5 second - timeout := 5 * time.Second haveResponse := false - if upstream.Timeout > 0 { - // convert upstream.Timeout(second) to nanosecond - timeout = time.Duration(upstream.Timeout) * time.Second - } proxy := httputil.NewSingleHostReverseProxy(remote) proxy.Director = nil @@ -60,12 +54,23 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s // record chat message from user record.Body = string(inBody) + requestBody, requestBodyOK := ParseRequestBody(inBody) + + // set timeout, default is 5 second + timeout := 5 * time.Second + if upstream.Timeout > 0 { + // convert upstream.Timeout(second) to nanosecond + timeout = time.Duration(upstream.Timeout) * time.Second + } + if requestBodyOK == nil && !requestBody.Stream { + timeout = 60 * time.Second + } // timeout out request go func() { time.Sleep(timeout) if !haveResponse { - log.Println("Timeout", upstream.Endpoint) + log.Println("Timeout upstream", upstream.Endpoint) errCtx = errors.New("timeout") cancel() } diff --git a/request_body.go b/request_body.go new file mode 100644 index 0000000..b92c610 --- /dev/null +++ b/request_body.go @@ -0,0 +1,32 @@ +package main + +import ( + "encoding/json" +) + +type Message struct { + Content string `json:"content"` + Role string `json:"role"` +} + +type RequestBody struct { + Model string `json:"model"` + Messages []Message `json:"messages"` + Stream bool `json:"stream"` + Temperature float64 `json:"temperature"` + TopP int64 `json:"top_p"` + PresencePenalty float64 `json:"presence_penalty"` + FrequencyPenalty float64 `json:"frequency_penalty"` +} + +func ParseRequestBody(data []byte) (RequestBody, error) { + ret := RequestBody{} + + var requestBody RequestBody + err := json.Unmarshal(data, &requestBody) + if err != nil { + return ret, err + } + + return requestBody, nil +}