add MasterAPITestSuite

Author: zhenghaoz

common/mock/openai.go

@@ -15,6 +15,8 @@
 package mock
 
 import (
+	"bytes"
+	"encoding/json"
 	"fmt"
 	"github.com/emicklei/go-restful/v3"
 	"github.com/sashabaranov/go-openai"
@@ -28,16 +30,15 @@ type OpenAIServer struct {
 	authToken  string
 	ready      chan struct{}
 
-	mockChatCompletion string
-	mockEmbeddings     []float32
+	mockEmbeddings []float32
 }
 
 func NewOpenAIServer() *OpenAIServer {
 	s := &OpenAIServer{}
 	ws := new(restful.WebService)
 	ws.Path("/v1").
-		Consumes(restful.MIME_XML, restful.MIME_JSON).
-		Produces(restful.MIME_JSON, restful.MIME_XML)
+		Consumes(restful.MIME_JSON).
+		Produces(restful.MIME_JSON, "text/event-stream")
 	ws.Route(ws.POST("chat/completions").
 		Reads(openai.ChatCompletionRequest{}).
 		Writes(openai.ChatCompletionResponse{}).
@@ -80,10 +81,6 @@ func (s *OpenAIServer) Close() error {
 	return s.httpServer.Close()
 }
 
-func (s *OpenAIServer) ChatCompletion(mock string) {
-	s.mockChatCompletion = mock
-}
-
 func (s *OpenAIServer) Embeddings(embeddings []float32) {
 	s.mockEmbeddings = embeddings
 }
@@ -95,13 +92,31 @@ func (s *OpenAIServer) chatCompletion(req *restful.Request, resp *restful.Respon
 		_ = resp.WriteError(http.StatusBadRequest, err)
 		return
 	}
-	_ = resp.WriteEntity(openai.ChatCompletionResponse{
-		Choices: []openai.ChatCompletionChoice{{
-			Message: openai.ChatCompletionMessage{
-				Content: s.mockChatCompletion,
-			},
-		}},
-	})
+	if r.Stream {
+		content := r.Messages[0].Content
+		for i := 0; i < len(content); i += 8 {
+			buf := bytes.NewBuffer(nil)
+			buf.WriteString("data: ")
+			encoder := json.NewEncoder(buf)
+			_ = encoder.Encode(openai.ChatCompletionStreamResponse{
+				Choices: []openai.ChatCompletionStreamChoice{{
+					Delta: openai.ChatCompletionStreamChoiceDelta{
+						Content: content[i:min(i+8, len(content))],
+					},
+				}},
+			})
+			buf.WriteString("\n")
+			_, _ = resp.Write(buf.Bytes())
+		}
+	} else {
+		_ = resp.WriteEntity(openai.ChatCompletionResponse{
+			Choices: []openai.ChatCompletionChoice{{
+				Message: openai.ChatCompletionMessage{
+					Content: r.Messages[0].Content,
+				},
+			}},
+		})
+	}
 }
 
 func (s *OpenAIServer) embeddings(req *restful.Request, resp *restful.Response) {

common/mock/openai_test.go

@@ -16,8 +16,11 @@ package mock
 
 import (
 	"context"
+	"github.com/juju/errors"
 	"github.com/sashabaranov/go-openai"
 	"github.com/stretchr/testify/suite"
+	"io"
+	"strings"
 	"testing"
 )
 
@@ -45,7 +48,6 @@ func (suite *OpenAITestSuite) TearDownSuite() {
 }
 
 func (suite *OpenAITestSuite) TestChatCompletion() {
-	suite.server.ChatCompletion("World")
 	resp, err := suite.client.CreateChatCompletion(
 		context.Background(),
 		openai.ChatCompletionRequest{
@@ -59,7 +61,39 @@ func (suite *OpenAITestSuite) TestChatCompletion() {
 		},
 	)
 	suite.NoError(err)
-	suite.Equal("World", resp.Choices[0].Message.Content)
+	suite.Equal("Hello", resp.Choices[0].Message.Content)
+}
+
+func (suite *OpenAITestSuite) TestChatCompletionStream() {
+	content := "In my younger and more vulnerable years my father gave me some advice that I've been turning over in" +
+		" my mind ever since. Whenever you feel like criticizing anyone, he told me, just remember that all the " +
+		"people in this world haven't had the advantages that you've had."
+	stream, err := suite.client.CreateChatCompletionStream(
+		context.Background(),
+		openai.ChatCompletionRequest{
+			Model: "qwen2.5",
+			Messages: []openai.ChatCompletionMessage{
+				{
+					Role:    openai.ChatMessageRoleUser,
+					Content: content,
+				},
+			},
+			Stream: true,
+		},
+	)
+	suite.NoError(err)
+	defer stream.Close()
+	var buffer strings.Builder
+	for {
+		var resp openai.ChatCompletionStreamResponse
+		resp, err = stream.Recv()
+		if errors.Is(err, io.EOF) {
+			suite.Equal(content, buffer.String())
+			return
+		}
+		suite.NoError(err)
+		buffer.WriteString(resp.Choices[0].Delta.Content)
+	}
 }
 
 func (suite *OpenAITestSuite) TestEmbeddings() {

master/master.go

@@ -29,6 +29,7 @@ import (
 	"github.com/emicklei/go-restful/v3"
 	"github.com/jellydator/ttlcache/v3"
 	"github.com/juju/errors"
+	"github.com/sashabaranov/go-openai"
 	"github.com/zhenghaoz/gorse/base"
 	"github.com/zhenghaoz/gorse/base/encoding"
 	"github.com/zhenghaoz/gorse/base/log"
@@ -71,6 +72,7 @@ type Master struct {
 	jobsScheduler  *task.JobsScheduler
 	cacheFile      string
 	managedMode    bool
+	openAIClient   *openai.Client
 
 	// cluster meta cache
 	metaStore meta.Database
@@ -116,6 +118,7 @@ type Master struct {
 // NewMaster creates a master node.
 func NewMaster(cfg *config.Config, cacheFile string, managedMode bool) *Master {
 	rand.Seed(time.Now().UnixNano())
+
 	// setup trace provider
 	tp, err := cfg.Tracing.NewTracerProvider()
 	if err != nil {
@@ -124,12 +127,18 @@ func NewMaster(cfg *config.Config, cacheFile string, managedMode bool) *Master {
 	otel.SetTracerProvider(tp)
 	otel.SetErrorHandler(log.GetErrorHandler())
 	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{}))
+
+	// setup OpenAI client
+	clientConfig := openai.DefaultConfig(cfg.OpenAI.AuthToken)
+	clientConfig.BaseURL = cfg.OpenAI.BaseURL
+
 	m := &Master{
 		// create task monitor
 		cacheFile:     cacheFile,
 		managedMode:   managedMode,
 		jobsScheduler: task.NewJobsScheduler(cfg.Master.NumJobs),
 		tracer:        progress.NewTracer("master"),
+		openAIClient:  openai.NewClientWithConfig(clientConfig),
 		// default ranking model
 		rankingModelName: "bpr",
 		rankingModelSearcher: ranking.NewModelSearcher(