Function streamHandler

ai-provider/local/executor_test.go:31–65

(c *gin.Context)

}

func streamHandler(c *gin.Context) {
	// 创建一个通道，用于监测客户端关闭连接的信号
	model := c.Query("model")
	p, err := PullModel(model, uuid.NewString(), nil)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	done := make(chan struct{})
	// 启动一个 goroutine 监听客户端关闭连接
	go func() {
		select {
		case <-c.Writer.CloseNotify():
			log.Info("client closed connection,close pipeline")
			taskExecutor.ClosePipeline(model, p.id)
		case <-done:
		}
	}()

	c.Stream(func(w io.Writer) bool {
		select {
		case msg, ok := <-p.channel:
			if !ok {
				return false
			}
			_, err := w.Write([]byte(fmt.Sprintf("%s\n", msg.Msg)))
			if err != nil {
				log.Error("write message error: %v", err)
				return false
			}
			return true
		}
	})
	done <- struct{}{}
}

func stopPull(c *gin.Context) {
	model := c.Query("model")