refactor: use goroutinue instead of mutex to sync messages

ai/model.go

@@ -22,17 +22,17 @@ type InvokeRequest struct {
 // InvokeResponse is the response for chat completions
 type InvokeResponse struct {
 	// Content is the content from llm api response
-	Content string
+	Content string `json:"content,omitempty"`
 	// ToolCalls is the toolCalls from llm api response
-	ToolCalls map[uint32][]*openai.ToolCall
+	ToolCalls map[uint32][]*openai.ToolCall `json:"tool_calls,omitempty"`
 	// ToolMessages is the tool messages from llm api response
-	ToolMessages []ToolMessage
+	ToolMessages []ToolMessage `json:"tool_messages,omitempty"`
 	// FinishReason is the finish reason from llm api response
-	FinishReason string
+	FinishReason string `json:"finish_reason,omitempty"`
 	// TokenUsage is the token usage from llm api response
-	TokenUsage TokenUsage
+	TokenUsage TokenUsage `json:"token_usage,omitempty"`
 	// AssistantMessage is the assistant message from llm api response, only present when finish reason is "tool_calls"
-	AssistantMessage interface{}
+	AssistantMessage interface{} `json:"assistant_message,omitempty"`
 }
 
 // TokenUsage is the token usage in Response
@@ -63,7 +63,7 @@ type ParameterProperty struct {
 type ToolMessage struct {
 	Role       string `json:"role"`
 	Content    string `json:"content"`
-	ToolCallId string `json:"tool_call_id"`
+	ToolCallID string `json:"tool_call_id"`
 }
 
 // ChainMessage is the message for chaining llm request with preceeding `tool_calls` response

ai/openai.go

@@ -8,6 +8,7 @@ import (
 	"github.com/yomorun/yomo/core/ylog"
 )
 
+// ConvertToInvokeResponse converts openai.ChatCompletionResponse struct to InvokeResponse struct.
 func ConvertToInvokeResponse(res *openai.ChatCompletionResponse, tcs map[uint32]openai.Tool) (*InvokeResponse, error) {
 	choice := res.Choices[0]
 	ylog.Debug(">>finish_reason", "reason", choice.FinishReason)

cli/serve.go

@@ -27,6 +27,7 @@ import (
 	"github.com/yomorun/yomo/pkg/log"
 
 	"github.com/yomorun/yomo/pkg/bridge/ai"
+	providerpkg "github.com/yomorun/yomo/pkg/bridge/ai/provider"
 	"github.com/yomorun/yomo/pkg/bridge/ai/provider/azopenai"
 	"github.com/yomorun/yomo/pkg/bridge/ai/provider/cfazure"
 	"github.com/yomorun/yomo/pkg/bridge/ai/provider/cfopenai"
@@ -78,7 +79,7 @@ var serveCmd = &cobra.Command{
 		}
 		if aiConfig != nil {
 			// add AI connection middleware
-			options = append(options, yomo.WithZipperConnMiddleware(ai.ConnMiddleware))
+			options = append(options, yomo.WithZipperConnMiddleware(ai.RegisterFunctionMW()))
 		}
 		// new zipper
 		zipper, err := yomo.NewZipper(
@@ -97,7 +98,7 @@ var serveCmd = &cobra.Command{
 			registerAIProvider(aiConfig)
 			// start the llm api server
 			go func() {
-				err := ai.Serve(aiConfig, listenAddr, fmt.Sprintf("token:%s", tokenString))
+				err := ai.Serve(aiConfig, listenAddr, fmt.Sprintf("token:%s", tokenString), ylog.Default())
 				if err != nil {
 					log.FailureStatusEvent(os.Stdout, err.Error())
 					return
@@ -118,37 +119,36 @@ func registerAIProvider(aiConfig *ai.Config) error {
 	for name, provider := range aiConfig.Providers {
 		switch name {
 		case "azopenai":
-			ai.RegisterProvider(azopenai.NewProvider(
+			providerpkg.RegisterProvider(azopenai.NewProvider(
 				provider["api_key"],
 				provider["api_endpoint"],
 				provider["deployment_id"],
 				provider["api_version"],
 			))
 		case "openai":
-			ai.RegisterProvider(openai.NewProvider(provider["api_key"], provider["model"]))
+			providerpkg.RegisterProvider(openai.NewProvider(provider["api_key"], provider["model"]))
 		case "cloudflare_azure":
-			ai.RegisterProvider(cfazure.NewProvider(
+			providerpkg.RegisterProvider(cfazure.NewProvider(
 				provider["endpoint"],
 				provider["api_key"],
 				provider["resource"],
 				provider["deployment_id"],
 				provider["api_version"],
 			))
 		case "cloudflare_openai":
-			ai.RegisterProvider(cfopenai.NewProvider(
+			providerpkg.RegisterProvider(cfopenai.NewProvider(
 				provider["endpoint"],
 				provider["api_key"],
 				provider["model"],
 			))
 		case "ollama":
-			ai.RegisterProvider(ollama.NewProvider(provider["api_endpoint"]))
+			providerpkg.RegisterProvider(ollama.NewProvider(provider["api_endpoint"]))
 		default:
 			log.WarningStatusEvent(os.Stdout, "unknown provider: %s", name)
 		}
 	}
 
-	// log.InfoStatusEvent(os.Stdout, "registered [%d] AI provider", len(ai.ListProviders()))
-	ylog.Info("registered AI providers", "len", len(ai.ListProviders()))
+	ylog.Info("registered AI providers", "len", len(providerpkg.ListProviders()))
 	return nil
 }
 

cli/test.go

@@ -214,7 +214,7 @@ var testPromptCmd = &cobra.Command{
 func getToolCallResult(tc *openai.ToolCall, tms []ai.ToolMessage) string {
 	result := ""
 	for _, tm := range tms {
-		if tm.ToolCallId == tc.ID {
+		if tm.ToolCallID == tc.ID {
 			result = tm.Content
 		}
 	}

core/connection.go

@@ -44,7 +44,8 @@ type Connection struct {
 	Logger          *slog.Logger
 }
 
-func newConnection(
+// NewConnection creates a new connection according to the parameters.
+func NewConnection(
 	id uint64,
 	name string, clientID string, clientType ClientType, md metadata.M, tags []uint32,
 	fconn frame.Conn, logger *slog.Logger,

core/connection_test.go

@@ -18,7 +18,7 @@ func TestConnection(t *testing.T) {
 		md       metadata.M
 	)
 
-	connection := newConnection(id, name, clientID, styp, md, observed, nil, ylog.Default())
+	connection := NewConnection(id, name, clientID, styp, md, observed, nil, ylog.Default())
 
 	t.Run("ConnectionInfo", func(t *testing.T) {
 		assert.Equal(t, id, connection.ID())

core/connector_test.go

@@ -117,5 +117,5 @@ func TestConnector(t *testing.T) {
 // mockConn returns a connection that only includes an ID and a name.
 // This function is used for unit testing purposes.
 func mockConn(id uint64, name string) *Connection {
-	return newConnection(id, name, "mock-id", ClientType(0), nil, []frame.Tag{0}, nil, ylog.Default())
+	return NewConnection(id, name, "mock-id", ClientType(0), nil, []frame.Tag{0}, nil, ylog.Default())
 }

core/router/router.go

@@ -33,9 +33,9 @@ type defaultRouter struct {
 	data map[frame.Tag]map[uint64]struct{}
 }
 
-// DefaultRouter provides a default implementation of `router`,
+// Default provides a default implementation of `router`,
 // It routes data according to observed tag and metadata.
-func Default() *defaultRouter {
+func Default() Router {
 	return &defaultRouter{
 		targets: make(map[uint64]string),
 		data:    make(map[frame.Tag]map[uint64]struct{}),

core/server.go

@@ -293,7 +293,7 @@ func (s *Server) createConnection(hf *frame.HandshakeFrame, md metadata.M, fconn
 	if hf.WantedTarget != "" {
 		md.Set(metadata.WantedTargetKey, hf.WantedTarget)
 	}
-	conn := newConnection(
+	conn := NewConnection(
 		incrID(),
 		hf.Name,
 		hf.ID,

options.go

@@ -112,6 +112,7 @@ var (
 		}
 	}
 
+	// WithVersionNegotiateFunc sets the version negotiate function for the zipper
 	WithVersionNegotiateFunc = func(f core.VersionNegotiateFunc) ZipperOption {
 		return func(zo *zipperOptions) {
 			zo.serverOption = append(zo.serverOption, core.WithVersionNegotiateFunc(f))

pkg/bridge/ai/ai.go

@@ -22,43 +22,48 @@
 	ErrConfigFormatError = errors.New("ai config format is incorrect")
 )
 
-// ConnMiddleware returns a ConnMiddleware that can be used to intercept the connection.
-func ConnMiddleware(next core.ConnHandler) core.ConnHandler {
-	return func(conn *core.Connection) {
-		connMd := conn.Metadata().Clone()
-		definition, ok := connMd.Get(ai.FunctionDefinitionKey)
+// RegisterFunctionMW returns a ConnMiddleware that can be used to register an ai function.
+func RegisterFunctionMW() core.ConnMiddleware { return registerFunction(register.GetRegister()) }
 
-		defer func() {
-			// definition does not be transmitted in mesh network, It only works for handshake.
-			conn.Metadata().Set(ai.FunctionDefinitionKey, "")
-			next(conn)
-			if ok {
-				register.UnregisterFunction(conn.ID(), connMd)
-				conn.Logger.Info("unregister ai function", "name", conn.Name(), "connID", conn.ID())
-			}
-		}()
+func registerFunction(r register.Register) core.ConnMiddleware {
+	return core.ConnMiddleware(func(next core.ConnHandler) core.ConnHandler {
+		return func(conn *core.Connection) {
+			connMd := conn.Metadata().Clone()
+			definition, ok := connMd.Get(ai.FunctionDefinitionKey)
 
-		if conn.ClientType() != core.ClientTypeStreamFunction || !ok {
-			return
-		}
+			defer func() {
+				if ok {
+					conn.Metadata().Set(ai.FunctionDefinitionKey, "")
+				}
+				// definition does not be transmitted in mesh network, It only works for handshake.
+				next(conn)
+				if ok {
+					register.UnregisterFunction(conn.ID(), connMd)
+					conn.Logger.Info("unregister ai function", "name", conn.Name(), "connID", conn.ID())
+				}
+			}()
 
-		for _, tag := range conn.ObserveDataTags() {
-			// register ai function
-			fd := ai.FunctionDefinition{}
-			err := json.Unmarshal([]byte(definition), &fd)
-			if err != nil {
-				conn.Logger.Error("unmarshal function definition", "error", err)
+			if conn.ClientType() != core.ClientTypeStreamFunction || !ok {
 				return
 			}
-			err = register.RegisterFunction(tag, &fd, conn.ID(), connMd)
-			if err != nil {
-				conn.Logger.Error("failed to register ai function", "name", conn.Name(), "tag", tag, "err", err)
-				return
+
+			for _, tag := range conn.ObserveDataTags() {
+				// register ai function
+				fd := ai.FunctionDefinition{}
+				err := json.Unmarshal([]byte(definition), &fd)
+				if err != nil {
+					conn.Logger.Error("unmarshal function definition", "error", err)
+					return
+				}
+				err = r.RegisterFunction(tag, &fd, conn.ID(), connMd)
+				if err != nil {
+					conn.Logger.Error("failed to register ai function", "name", conn.Name(), "tag", tag, "err", err)
+					return
+				}
+				conn.Logger.Info("register ai function success", "name", conn.Name(), "tag", tag, "definition", string(definition))
 			}
-			conn.Logger.Info("register ai function success", "name", conn.Name(), "tag", tag, "definition", string(definition))
 		}
-
-	}
+	})
 }
 
 // Config is the configuration of AI bridge.