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

generate_response.go

@@ -1,6 +1,7 @@
 package ollama
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/json"
 	"errors"
@@ -16,6 +17,7 @@ type GenerateResponseRequest struct {
 	Format      string                          `json:"format,omitempty"`
 	System      string                          `json:"system,omitempty"`
 	Stream      *bool                           `json:"stream,omitempty"`
+	Think       *bool                           `json:"think,omitempty"`
 	Raw         *bool                           `json:"raw,omitempty"`
 	KeepAlive   string                          `json:"keep_alive,omitempty"`
 	Options     *GenerateResponseRequestOptions `json:"options,omitempty"`
@@ -59,15 +61,32 @@ type GenerateResponseResponse struct {
 	} `json:"logprobs"`
 }
 
-func (o Ollama) GenerateResponse(reqBody GenerateResponseRequest) (GenerateResponseResponse, int, error) {
+type GenerateResponseResponseStream struct {
+	Model              string `json:"model"`
+	CreatedAt          string `json:"created_at"`
+	Response           string `json:"response"`
+	Thinking           string `json:"thinking"`
+	Done               bool   `json:"done"`
+	DoneReason         string `json:"done_reason"`
+	TotalDuration      int    `json:"total_duration"`
+	LoadDuration       int    `json:"load_duration"`
+	PromptEvalCount    int    `json:"prompt_eval_count"`
+	PromptEvalDuration int    `json:"prompt_eval_duration"`
+	EvalCount          int    `json:"eval_count"`
+	EvalDuration       int    `json:"eval_duration"`
+}
+
+func (o Ollama) GenerateResponse(reqBody GenerateResponseRequest) (GenerateResponseResponse, error) {
+	reqBody.Stream = PtrOf(false)
+
 	reqBodyBytes, err := json.Marshal(reqBody)
 	if err != nil {
-		return GenerateResponseResponse{}, -1, err
+		return GenerateResponseResponse{}, err
 	}
 
 	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/generate", o.baseUrl), bytes.NewReader(reqBodyBytes))
 	if err != nil {
-		return GenerateResponseResponse{}, -1, err
+		return GenerateResponseResponse{}, err
 	}
 
 	for key, val := range o.customHeaders {
@@ -77,17 +96,68 @@ func (o Ollama) GenerateResponse(reqBody GenerateResponseRequest) (GenerateRespo
 
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
-		return GenerateResponseResponse{}, -1, err
+		return GenerateResponseResponse{}, err
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != 200 {
-		return GenerateResponseResponse{}, resp.StatusCode, errors.New("status code is not 200")
+	if resp.StatusCode != http.StatusOK {
+		return GenerateResponseResponse{}, errors.New("status code is not 200")
 	}
 
 	var respBody GenerateResponseResponse
 	if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
-		return GenerateResponseResponse{}, -1, err
+		return GenerateResponseResponse{}, err
 	}
-	return respBody, resp.StatusCode, nil
+	return respBody, nil
+}
+
+func (o Ollama) GenerateResponseStream(reqBody GenerateResponseRequest, onChunk func(chunk GenerateResponseResponseStream)) error {
+	reqBody.Stream = PtrOf(true)
+
+	reqBodyBytes, err := json.Marshal(reqBody)
+	if err != nil {
+		return err
+	}
+
+	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/generate", 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 GenerateResponseResponseStream
+		if err := json.Unmarshal(line, &chunk); err != nil {
+			return err
+		}
+
+		onChunk(chunk)
+		if chunk.Done {
+			return nil
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+
+	return nil
 }