feat(ai-bridge): handle llm-sfn return nothing or multiple results returned (#833)

Author: woorui

core/serverless/context_llm.go

@@ -9,8 +9,7 @@ import (
 
 // ReadLLMArguments reads LLM function arguments
 func (c *Context) ReadLLMArguments(args any) error {
-	fnCall := &ai.FunctionCall{}
-	err := fnCall.FromBytes(c.data)
+	fnCall, err := c.LLMFunctionCall()
 	if err != nil {
 		return err
 	}
@@ -27,13 +26,17 @@ func (c *Context) WriteLLMResult(result string) error {
 	if c.fnCall == nil {
 		return errors.New("no function call, can't write result")
 	}
+	if c.fnCall.IsOK && c.fnCall.Result != "" {
+		return errors.New("LLM function can only be called once")
+	}
 	// function call
 	c.fnCall.IsOK = true
 	c.fnCall.Result = result
 	buf, err := c.fnCall.Bytes()
 	if err != nil {
 		return err
 	}
+	c.data = buf
 	return c.Write(ai.ReducerTag, buf)
 }
 

pkg/bridge/ai/service.go

@@ -427,7 +427,7 @@ func (s *Service) GetChatCompletions(ctx context.Context, req openai.ChatComplet
 		}
 		promptUsage = resp.Usage.PromptTokens
 		completionUsage = resp.Usage.CompletionTokens
-		totalUsage = resp.Usage.CompletionTokens
+		totalUsage = resp.Usage.TotalTokens
 
 		ylog.Debug(" #1 first call", "response", fmt.Sprintf("%+v", resp))
 		// it is a function call

serverless/mock/mock_context.go

@@ -7,7 +7,6 @@ import (
 
 	"github.com/yomorun/yomo/ai"
 	"github.com/yomorun/yomo/serverless"
-	"github.com/yomorun/yomo/serverless/guest"
 )
 
 var _ serverless.Context = (*MockContext)(nil)
@@ -55,7 +54,7 @@ func (c *MockContext) Metadata(_ string) (string, bool) {
 
 // HTTP returns the HTTP interface.H
 func (m *MockContext) HTTP() serverless.HTTP {
-	return &guest.GuestHTTP{}
+	panic("not implemented, to use `net/http` package")
 }
 
 // Write writes the data with the given tag.
@@ -106,8 +105,13 @@ func (c *MockContext) WriteLLMResult(result string) error {
 	defer c.mu.Unlock()
 
 	if c.fnCall == nil {
-		return errors.New("no function call, can't write result")
+		fnCall, err := c.LLMFunctionCall()
+		if err != nil {
+			return err
+		}
+		c.fnCall = fnCall
 	}
+
 	// function call
 	c.fnCall.IsOK = true
 	c.fnCall.Result = result

serverless/mock/mock_context_test.go

@@ -89,3 +89,14 @@ func TestWriteLLMResult(t *testing.T) {
 	assert.Equal(t, ai.ReducerTag, res[0].Tag)
 	assert.Equal(t, jsonStrWithResult("test result"), string(res[0].Data))
 }
+
+func TestWriteLLMResultWithoutRead(t *testing.T) {
+	ctx := NewMockContext([]byte(jsonStr), 0x10)
+	// write
+	err := ctx.WriteLLMResult("test result")
+	assert.NoError(t, err)
+
+	res := ctx.RecordsWritten()
+	assert.Equal(t, ai.ReducerTag, res[0].Tag)
+	assert.Equal(t, jsonStrWithResult("test result"), string(res[0].Data))
+}

sfn.go

@@ -3,6 +3,7 @@ package yomo
 import (
 	"context"
 	"errors"
+	"log/slog"
 
 	"github.com/robfig/cron/v3"
 
@@ -12,6 +13,7 @@ import (
 	"github.com/yomorun/yomo/core/serverless"
 	"github.com/yomorun/yomo/pkg/id"
 	"github.com/yomorun/yomo/pkg/trace"
+	yserverless "github.com/yomorun/yomo/serverless"
 	"go.opentelemetry.io/otel/attribute"
 )
 
@@ -251,6 +253,7 @@ func (s *streamFunction) onDataFrame(dataFrame *frame.DataFrame) {
 
 			serverlessCtx := serverless.NewContext(s.client, dataFrame.Tag, md, dataFrame.Payload)
 			s.fn(serverlessCtx)
+			checkLLMFunctionCall(s.client.Logger, serverlessCtx)
 		}(dataFrame)
 	} else if s.pfn != nil {
 		data := dataFrame.Payload
@@ -270,3 +273,14 @@ func (s *streamFunction) SetErrorHandler(fn func(err error)) {
 func (s *streamFunction) Init(fn func() error) error {
 	return fn()
 }
+
+func checkLLMFunctionCall(logger *slog.Logger, serverlessCtx yserverless.Context) {
+	fc, err := serverlessCtx.LLMFunctionCall()
+	if err != nil {
+		// it's not a LLM function call ctx
+		return
+	}
+	if !fc.IsOK {
+		logger.Warn("The function is not returning anything， please check if `WriteLLMResult()` has been called")
+	}
+}