Fix in server

server/lm/cmd.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os/exec"
+	"strings"
 
 	"github.com/labstack/echo/v4"
 	"github.com/synw/agent-smith/server/state"
@@ -51,11 +52,13 @@ func RunCmd(
 	enc := json.NewEncoder(c.Response())
 	enc.SetEscapeHTML(false)
 	i := 0
+	buf := []string{}
 	for scanner.Scan() {
 		i++
 		token := scanner.Text()
+		buf = append(buf, token)
 		if state.IsVerbose {
-			fmt.Print(token)
+			go fmt.Print(token)
 		}
 		StreamMsg(createMsg(token, i), c, enc)
 	}
@@ -71,5 +74,17 @@ func RunCmd(
 	if err := cmd.Wait(); err != nil {
 		msg := fmt.Errorf("Command finished with error: %v", err)
 		errCh <- createErrorMsg(msg.Error())
+	} else {
+		//fmt.Println("Command finished")
+		res := make(map[string]interface{})
+		res["text"] = strings.Join(buf, "")
+		endmsg := types.StreamedMessage{
+			Num:     0,
+			Content: "result",
+			MsgType: types.SystemMsgType,
+			Data:    res,
+		}
+		StreamMsg(endmsg, c, enc)
+		ch <- endmsg
 	}
 }

server/lm/oai.go

@@ -111,8 +111,10 @@ func OaiInfer(
 	ntokens := 0
 	for {
 		response, err := stream.Recv()
+		fmt.Println("SERR", err)
+		fmt.Println("RESP", response)
 		if errors.Is(err, io.EOF) {
-			//fmt.Println("\nStream finished")
+			fmt.Println("\nStream EOF")
 			emittingElapsed := time.Since(startEmitting)
 			if state.IsVerbose {
 				fmt.Println("\n\nEmitting time:", emittingElapsed)
@@ -155,7 +157,7 @@ func OaiInfer(
 			}
 			StreamMsg(endmsg, c, enc)
 			ch <- endmsg
-			return nil
+			break
 		}
 		if err != nil {
 			fmt.Printf("\nStream error: %v\n", err)
@@ -203,4 +205,5 @@ func OaiInfer(
 		}
 		ntokens++
 	}
+	return nil
 }