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

create_model.go

@@ -1,6 +1,7 @@
 package ollama
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/json"
 	"errors"
@@ -41,15 +42,24 @@ type CreateModelResponse struct {
 	Status string `json:"status"`
 }
 
-func (o Ollama) CreateModel(reqBody CreateModelRequest) (CreateModelResponse, int, error) {
+type CreateModelResponseStream struct {
+	Status    string `json:"status"`
+	Digest    string `json:"digest"`
+	Total     int    `json:"total"`
+	Completed int    `json:"completed"`
+}
+
+func (o Ollama) CreateModel(reqBody CreateModelRequest) (CreateModelResponse, error) {
+	reqBody.Stream = PtrOf(false)
+
 	reqBodyBytes, err := json.Marshal(reqBody)
 	if err != nil {
-		return CreateModelResponse{}, -1, err
+		return CreateModelResponse{}, err
 	}
 
 	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/create", o.baseUrl), bytes.NewReader(reqBodyBytes))
 	if err != nil {
-		return CreateModelResponse{}, -1, err
+		return CreateModelResponse{}, err
 	}
 
 	for key, val := range o.customHeaders {
@@ -59,17 +69,68 @@ func (o Ollama) CreateModel(reqBody CreateModelRequest) (CreateModelResponse, in
 
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
-		return CreateModelResponse{}, -1, err
+		return CreateModelResponse{}, err
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != 200 {
-		return CreateModelResponse{}, resp.StatusCode, errors.New("status code is not 200")
+	if resp.StatusCode != http.StatusOK {
+		return CreateModelResponse{}, errors.New("status code is not 200")
 	}
 
 	var respBody CreateModelResponse
 	if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
-		return CreateModelResponse{}, -1, err
+		return CreateModelResponse{}, err
 	}
-	return respBody, resp.StatusCode, nil
+	return respBody, nil
+}
+
+func (o Ollama) CreateModelStream(reqBody CreateModelRequest, onChunk func(chunk CreateModelResponseStream)) error {
+	reqBody.Stream = PtrOf(true)
+
+	reqBodyBytes, err := json.Marshal(reqBody)
+	if err != nil {
+		return err
+	}
+
+	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/create", 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 CreateModelResponseStream
+		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
 }