diff --git a/relay/adaptor/ollama/main.go b/relay/adaptor/ollama/main.go index 6a1d334d1a..43317ff66f 100644 --- a/relay/adaptor/ollama/main.go +++ b/relay/adaptor/ollama/main.go @@ -31,6 +31,8 @@ func ConvertRequest(request model.GeneralOpenAIRequest) *ChatRequest { TopP: request.TopP, FrequencyPenalty: request.FrequencyPenalty, PresencePenalty: request.PresencePenalty, + NumPredict: request.MaxTokens, + NumCtx: request.NumCtx, }, Stream: request.Stream, } @@ -118,8 +120,10 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC common.SetEventStreamHeaders(c) for scanner.Scan() { - data := strings.TrimPrefix(scanner.Text(), "}") - data = data + "}" + data := scanner.Text() + if strings.HasPrefix(data, "}") { + data = strings.TrimPrefix(data, "}") + "}" + } var ollamaResponse ChatResponse err := json.Unmarshal([]byte(data), &ollamaResponse) diff --git a/relay/adaptor/ollama/model.go b/relay/adaptor/ollama/model.go index 29430e1c7c..7039984fcc 100644 --- a/relay/adaptor/ollama/model.go +++ b/relay/adaptor/ollama/model.go @@ -7,6 +7,8 @@ type Options struct { TopP float64 `json:"top_p,omitempty"` FrequencyPenalty float64 `json:"frequency_penalty,omitempty"` PresencePenalty float64 `json:"presence_penalty,omitempty"` + NumPredict int `json:"num_predict,omitempty"` + NumCtx int `json:"num_ctx,omitempty"` } type Message struct { diff --git a/relay/model/general.go b/relay/model/general.go index 229a61c160..c34c1c2d5d 100644 --- a/relay/model/general.go +++ b/relay/model/general.go @@ -29,6 +29,7 @@ type GeneralOpenAIRequest struct { Dimensions int `json:"dimensions,omitempty"` Instruction string `json:"instruction,omitempty"` Size string `json:"size,omitempty"` + NumCtx int `json:"num_ctx,omitempty"` } func (r GeneralOpenAIRequest) ParseInput() []string {