go-ollama: changed development to dev container and added streaming capabilities to methods v0.1.1

push_model.go

@@ -1,6 +1,7 @@
 package ollama
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/json"
 	"errors"
@@ -18,15 +19,23 @@ type PushModelResponse struct {
 	Status string `json:"status"`
 }
 
-func (o Ollama) PushModel(reqBody PushModelRequest) (PushModelResponse, int, error) {
+type PushModelResponseStream struct {
+	Status    string `json:"status"`
+	Digest    string `json:"digest"`
+	Total     int    `json:"total"`
+	Completed int    `json:"completed"`
+}
+
+func (o Ollama) PushModel(reqBody PushModelRequest) (PushModelResponse, error) {
+	reqBody.Stream = PtrOf(false)
 	reqBodyBytes, err := json.Marshal(reqBody)
 	if err != nil {
-		return PushModelResponse{}, -1, err
+		return PushModelResponse{}, err
 	}
 
-	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/pull", o.baseUrl), bytes.NewReader(reqBodyBytes))
+	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/push", o.baseUrl), bytes.NewReader(reqBodyBytes))
 	if err != nil {
-		return PushModelResponse{}, -1, err
+		return PushModelResponse{}, err
 	}
 
 	for key, val := range o.customHeaders {
@@ -36,17 +45,68 @@ func (o Ollama) PushModel(reqBody PushModelRequest) (PushModelResponse, int, err
 
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
-		return PushModelResponse{}, -1, err
+		return PushModelResponse{}, err
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != 200 {
-		return PushModelResponse{}, resp.StatusCode, errors.New("status code is not 200")
+	if resp.StatusCode != http.StatusOK {
+		return PushModelResponse{}, errors.New("status code is not 200")
 	}
 
 	var respBody PushModelResponse
 	if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
-		return PushModelResponse{}, -1, err
+		return PushModelResponse{}, err
 	}
-	return respBody, resp.StatusCode, nil
+	return respBody, nil
+}
+
+func (o Ollama) PushModelStream(reqBody PushModelRequest, onChunk func(chunk PushModelResponseStream)) error {
+	reqBody.Stream = PtrOf(true)
+
+	reqBodyBytes, err := json.Marshal(reqBody)
+	if err != nil {
+		return err
+	}
+
+	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/push", o.baseUrl), bytes.NewReader(reqBodyBytes))
+	if err != nil {
+		return err
+	}
+
+	for key, val := range o.customHeaders {
+		req.Header.Set(key, val)
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return errors.New("status code is not 200")
+	}
+
+	scanner := bufio.NewScanner(resp.Body)
+
+	for scanner.Scan() {
+		line := bytes.TrimSpace(scanner.Bytes())
+
+		var chunk PushModelResponseStream
+		if err := json.Unmarshal(line, &chunk); err != nil {
+			return err
+		}
+
+		onChunk(chunk)
+		if chunk.Status == "success" {
+			break
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+
+	return nil
 }