Add streaming for code completion between a prompt and a suffix

Author: xyproto

v2/code.go

@@ -58,8 +58,7 @@ func (oc *Config) GetBetweenResponse(prompt, suffix string) (OutputResponse, err
 			return res, nil
 		}
 	}
-	var reqBody GenerateRequest
-	reqBody = GenerateRequest{
+	reqBody := GenerateRequest{
 		Model:  oc.ModelName,
 		Prompt: prompt,
 		Suffix: suffix,
@@ -125,3 +124,71 @@ func (oc *Config) Complete(codeStart, codeEnd string) (string, error) {
 	}
 	return response.Response, nil
 }
+
+// StreamBetween sends a request to the Ollama API and returns the generated output via a callback function.
+// The callback function is given a string and "true" when the streaming is done (or if an error occurred).
+func (oc *Config) StreamBetween(callbackFunction func(string, bool), prompt, suffix string) error {
+	defer callbackFunction("", true)
+	var (
+		temperature float64
+		seed        = oc.SeedOrNegative
+	)
+	if prompt == "" {
+		return errors.New("the prompt can not be empty")
+	}
+	if seed < 0 {
+		temperature = oc.TemperatureIfNegativeSeed
+	}
+	reqBody := GenerateRequest{
+		Model:  oc.ModelName,
+		System: oc.SystemPrompt,
+		Prompt: prompt,
+		Suffix: suffix,
+		Stream: true,
+		Options: RequestOptions{
+			Seed:        seed,        // set to -1 to make it random
+			Temperature: temperature, // set to 0 together with a specific seed to make output reproducible
+		},
+	}
+	if oc.ContextLength != 0 {
+		reqBody.Options.ContextLength = oc.ContextLength
+	}
+	reqBytes, err := json.Marshal(reqBody)
+	if err != nil {
+		return err
+	}
+	if oc.Verbose {
+		fmt.Printf("Sending request to %s/api/generate: %s\n", oc.ServerAddr, string(reqBytes))
+	}
+	HTTPClient := &http.Client{
+		Timeout: oc.HTTPTimeout,
+	}
+	resp, err := HTTPClient.Post(oc.ServerAddr+"/api/generate", mimeJSON, bytes.NewBuffer(reqBytes))
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	var (
+		decoder = json.NewDecoder(resp.Body)
+		first   = true
+		answer  string
+	)
+	for {
+		var genResp GenerateResponse
+		if err := decoder.Decode(&genResp); err != nil {
+			break
+		}
+		answer = genResp.Response
+		if first {
+			if len(answer) > 0 && answer[0] == ' ' {
+				answer = answer[1:]
+			}
+			first = false
+		}
+		callbackFunction(answer, false)
+		if genResp.Done {
+			break
+		}
+	}
+	return nil
+}

v2/code_test.go

@@ -11,12 +11,15 @@ func TestBetween(t *testing.T) {
 	const (
 		codeStart = "def compute_gcd(a, b):"
 		codeEnd   = "    return result"
+		verbose   = true
 	)
+
 	oc := New(codeCompleteModel)
-	oc.Verbose = true
+	oc.Verbose = verbose
 	if err := oc.PullIfNeeded(true); err != nil {
 		t.Fatalf("Failed to pull model: %v", err)
 	}
+
 	response, err := oc.GetBetweenResponse(codeStart, codeEnd)
 	if err != nil {
 		t.Fatalf("Failed to get code completion: %v", err)
@@ -30,11 +33,42 @@ func TestCodeCompletion(t *testing.T) {
 		codeEnd   = "    return result"
 		verbose   = true
 	)
+
 	oc := New(codeCompleteModel)
-	oc.Verbose = true
+	oc.Verbose = verbose
+
 	codeBetween, err := oc.Complete(codeStart, codeEnd)
 	if err != nil {
 		t.Fatal(err)
 	}
 	fmt.Printf("%s\n%s\n%s\n", codeStart, codeBetween, codeEnd)
 }
+
+func TestStreamBetween(t *testing.T) {
+	const (
+		codeStart = "def compute_gcd(a, b):"
+		codeEnd   = "    return result"
+		verbose   = false
+	)
+
+	oc := New(codeCompleteModel)
+	oc.Verbose = verbose
+
+	err := oc.PullIfNeeded(true)
+	if err != nil {
+		t.Fatalf("Failed to pull model: %v", err)
+	}
+
+	callbackFunction := func(partialResult string, streamingDone bool) {
+		if !streamingDone {
+			fmt.Printf("%s", partialResult)
+		} else {
+			fmt.Println(codeEnd)
+		}
+	}
+
+	fmt.Print(codeStart)
+	if err := oc.StreamBetween(callbackFunction, codeStart, codeEnd); err != nil {
+		t.Fatalf("Failed to get streamed : %v", err)
+	}
+}