Fix /chat/completion response in echo mode (#362)

* Fix /chat/completion response in echo mode + update tests accordingly

Signed-off-by: Maya Barnea <mayab@il.ibm.com>

* Fix echo mode for grpc + lint

Signed-off-by: Maya Barnea <mayab@il.ibm.com>

* Fix echo in grpc

Signed-off-by: Maya Barnea <mayab@il.ibm.com>

---------

Signed-off-by: Maya Barnea <mayab@il.ibm.com>

Author: mayabar

pkg/dataset/dataset.go

@@ -57,9 +57,9 @@ type EchoDataset struct{}
 // if max-tokens is defined in the request and response's length is >= it value, finish reason is set to LENGTH,
 // otherwise finish reason is STOP
 func (ed *EchoDataset) GetResponseTokens(req openaiserverapi.Request) (*openaiserverapi.Tokenized, string, error) {
-	tokens := req.TokenizedPrompt()
+	tokens := req.TokenizedEchoResponse()
 	maxTokens := req.GetMaxCompletionTokens()
-	return tokens, common.FinishReason(maxTokens, len(tokens.Tokens)), nil
+	return tokens, common.FinishReason(maxTokens, tokens.Length()), nil
 }
 
 func (ed *EchoDataset) Close() error {

pkg/dataset/dataset_test.go

@@ -187,6 +187,7 @@ var _ = Describe("Echo Dataset", Ordered, func() {
 				Prompt: theText,
 			}
 			req.SetTokenizedPrompt(&openaiserverapi.Tokenized{Tokens: tokens, Strings: strTokens})
+			req.SetTokenizedEchoResponse(&openaiserverapi.Tokenized{Tokens: tokens, Strings: strTokens})
 			tokens, finishReason, err := dataset.GetResponseTokens(req)
 			Expect(err).ShouldNot(HaveOccurred())
 			Expect(tokens.Strings).Should(Equal(strTokens))
@@ -199,6 +200,7 @@ var _ = Describe("Echo Dataset", Ordered, func() {
 				MaxTokens: &maxTokens,
 			}
 			req.SetTokenizedPrompt(&openaiserverapi.Tokenized{Tokens: tokens, Strings: strTokens})
+			req.SetTokenizedEchoResponse(&openaiserverapi.Tokenized{Tokens: tokens, Strings: strTokens})
 
 			tokens, finishReason, err := dataset.GetResponseTokens(req)
 			Expect(err).ShouldNot(HaveOccurred())
@@ -212,12 +214,34 @@ var _ = Describe("Echo Dataset", Ordered, func() {
 				MaxTokens: &maxTokens,
 			}
 			req.SetTokenizedPrompt(&openaiserverapi.Tokenized{Tokens: tokens, Strings: strTokens})
+			req.SetTokenizedEchoResponse(&openaiserverapi.Tokenized{Tokens: tokens, Strings: strTokens})
 
 			tokens, finishReason, err := dataset.GetResponseTokens(req)
 			Expect(err).ShouldNot(HaveOccurred())
 			Expect(tokens.Strings).Should(Equal(strTokens))
 			Expect(finishReason).Should(Equal(common.LengthFinishReason))
 		})
+		It("should return the last message in chat completion", func() {
+			req := &openaiserverapi.ChatCompletionRequest{
+				Messages: []openaiserverapi.Message{
+					{Role: openaiserverapi.RoleUser, Content: openaiserverapi.Content{Raw: "user message1"}},
+					{Role: openaiserverapi.RoleAssistant, Content: openaiserverapi.Content{Raw: "assistant message1"}},
+					{Role: openaiserverapi.RoleUser, Content: openaiserverapi.Content{Raw: testPrompt}},
+				},
+			}
+			promptTokens, promptStrTokens, err := tokenizer.Encode(req.GetFullPrompt(), "")
+			Expect(err).ShouldNot(HaveOccurred())
+			respTokens, resptStrTokens, err := tokenizer.Encode(testPrompt, "")
+			Expect(err).ShouldNot(HaveOccurred())
+
+			req.SetTokenizedPrompt(&openaiserverapi.Tokenized{Tokens: promptTokens, Strings: promptStrTokens})
+			req.SetTokenizedEchoResponse(&openaiserverapi.Tokenized{Tokens: respTokens, Strings: resptStrTokens})
+
+			tokens, _, err := dataset.GetResponseTokens(req)
+			Expect(err).ShouldNot(HaveOccurred())
+
+			Expect(tokens.Strings).Should(Equal(resptStrTokens))
+		})
 	})
 
 	DescribeTable("should work correctly in echo mode",
@@ -235,6 +259,7 @@ var _ = Describe("Echo Dataset", Ordered, func() {
 				req = &textReq
 			}
 			req.SetTokenizedPrompt(&openaiserverapi.Tokenized{Tokens: promptTokens, Strings: promptStrTokens})
+			req.SetTokenizedEchoResponse(&openaiserverapi.Tokenized{Tokens: promptTokens, Strings: promptStrTokens})
 
 			tokens, finishReason, err := dataset.GetResponseTokens(req)
 			Expect(err).NotTo(HaveOccurred())

pkg/llm-d-inference-sim/chat_completion.go

@@ -77,6 +77,14 @@ func (c *chatCompletionRequest) createResponseContext(reqCtx requestContext, dis
 	}
 }
 
+func (c *chatCompletionReqCtx) getEchoTokens() ([]uint32, []string, error) {
+	lastMsg := ""
+	if len(c.req.Messages) > 0 {
+		lastMsg = c.req.Messages[len(c.req.Messages)-1].Content.Raw
+	}
+	return c.sim.tokenizer.Encode(lastMsg, "")
+}
+
 var _ request = (*chatCompletionRequest)(nil)
 
 // Implementation of requestContext for /chat/completions requests

pkg/llm-d-inference-sim/generation.go

@@ -72,6 +72,14 @@ func (g *generationReqCtx) request() request {
 	return g.req
 }
 
+func (g *generationReqCtx) getEchoTokens() ([]uint32, []string, error) {
+	tokenisedResponse := g.req.TokenizedEchoResponse()
+	if tokenisedResponse != nil {
+		return tokenisedResponse.Tokens, tokenisedResponse.Strings, nil
+	}
+	return g.sim.tokenizer.Encode(g.req.Prompt, "")
+}
+
 func (g *generationReqCtx) kvCacheOnRequestStart() (hitRate float64, oaiServerError *openaiserverapi.Error) {
 	if g.sim.config.EnableKVCache {
 		var err error

pkg/llm-d-inference-sim/grpc.go

@@ -144,6 +144,7 @@ func (s *VllmSimulator) pbRequestToRequest(in *pb.GenerateRequest) *generationRe
 		prompt := &openaiserverapi.Tokenized{}
 		prompt.Tokens = in.GetTokenized().InputIds
 		req.SetTokenizedPrompt(prompt)
+		req.SetTokenizedEchoResponse(prompt)
 	} else {
 		req.Prompt = in.GetText()
 	}

pkg/llm-d-inference-sim/request.go

@@ -48,6 +48,7 @@ type requestContext interface {
 	createToolCalls() ([]openaiserverapi.ToolCall, int, string, error)
 	handleRequest() (responseContext, *openaiserverapi.Error)
 	responseChannel() chan *responseInfo
+	getEchoTokens() ([]uint32, []string, error)
 }
 
 type baseRequestContext struct {
@@ -92,6 +93,21 @@ func (b *baseRequestContext) tokenize() *openaiserverapi.Error {
 		Tokens:  tokens,
 		Strings: textTokens,
 	})
+
+	if b.sim.config.Mode == common.ModeEcho {
+		tokens, textTokens, err = b.getEchoTokens()
+		if err != nil {
+			b.sim.logger.Error(err, "failed to tokenize echo mode response")
+			serverErr := openaiserverapi.NewError("Failed to tokenize  echo mode response, "+err.Error(), fasthttp.StatusInternalServerError, nil)
+			return &serverErr
+		}
+
+		req.SetTokenizedEchoResponse(&openaiserverapi.Tokenized{
+			Tokens:  tokens,
+			Strings: textTokens,
+		})
+	}
+
 	return nil
 }
 

pkg/llm-d-inference-sim/text_completion.go

@@ -101,6 +101,10 @@ func (t *textCompletionReqCtx) createToolCalls() ([]openaiserverapi.ToolCall, in
 	return nil, 0, "", nil
 }
 
+func (t *textCompletionReqCtx) getEchoTokens() ([]uint32, []string, error) {
+	return t.sim.tokenizer.Encode(t.req.Prompt, "")
+}
+
 var _ requestContext = (*textCompletionReqCtx)(nil)
 
 // Implementation of responseContext for /completions requests

pkg/openai-server-api/request.go

@@ -80,6 +80,10 @@ type Request interface {
 	TokenizedPrompt() *Tokenized
 	// SetTokenizedPrompt sets the tokenized prompt
 	SetTokenizedPrompt(tokenized *Tokenized)
+	// TokenizedEchoResponse returns the tokenized response in echo mode
+	TokenizedEchoResponse() *Tokenized
+	// SetTokenizedEchoResponse sets the tokenized response in echo mode
+	SetTokenizedEchoResponse(tokenized *Tokenized)
 	// CacheThresholdFinishReason returns cacheThresholdFinishReason,  when true,
 	// forces a cache_threshold finish reason
 	CacheThresholdFinishReason() bool
@@ -112,6 +116,8 @@ type baseCompletionRequest struct {
 	cacheThresholdFinishReason bool
 	// tokenizedPrompt is the tokenized prompt
 	tokenizedPrompt *Tokenized
+	// tokenizedEchoResponse is the tokenized response in echo mode, exists only in echo mode
+	tokenizedEchoResponse *Tokenized
 }
 
 type KVTransferParams struct {
@@ -245,6 +251,16 @@ func (b *baseCompletionRequest) SetTokenizedPrompt(tokenized *Tokenized) {
 	b.tokenizedPrompt = tokenized
 }
 
+// TokenizedEchoResponse returns the tokenized response in echo mode
+func (b *baseCompletionRequest) TokenizedEchoResponse() *Tokenized {
+	return b.tokenizedEchoResponse
+}
+
+// SetTokenizedEchoResponse sets the tokenized response in echo mode
+func (b *baseCompletionRequest) SetTokenizedEchoResponse(tokenized *Tokenized) {
+	b.tokenizedEchoResponse = tokenized
+}
+
 // ChatCompletionRequest defines structure of /chat/completion request
 type ChatCompletionRequest struct {
 	baseCompletionRequest