feat: add credential for the build function for source and reducer (#…

…1001) This allows developers to create the builder for the source or reducer with the credential.
yomorun · Feb 7, 2025 · 825dd41 · 825dd41
1 parent bb7babb
commit 825dd41
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 12 deletions.
diff --git a/pkg/bridge/ai/api_server.go b/pkg/bridge/ai/api_server.go
@@ -83,8 +83,8 @@ func NewBasicAPIServer(config *Config, provider provider.LLMProvider, source yom
 
 	opts := &ServiceOptions{
 		Logger:         logger,
-		SourceBuilder:  func() yomo.Source { return source },
-		ReducerBuilder: func() yomo.StreamFunction { return reducer },
+		SourceBuilder:  func(_ string) yomo.Source { return source },
+		ReducerBuilder: func(_ string) yomo.StreamFunction { return reducer },
 	}
 	service := NewService(provider, opts)
 

diff --git a/pkg/bridge/ai/api_server_test.go b/pkg/bridge/ai/api_server_test.go
@@ -47,8 +47,8 @@ func TestServer(t *testing.T) {
 	}
 
 	service := newService(pd, newCaller, &ServiceOptions{
-		SourceBuilder:     func() yomo.Source { return flow },
-		ReducerBuilder:    func() yomo.StreamFunction { return flow },
+		SourceBuilder:     func(_ string) yomo.Source { return flow },
+		ReducerBuilder:    func(_ string) yomo.StreamFunction { return flow },
 		MetadataExchanger: func(_ string) (metadata.M, error) { return metadata.M{"hello": "llm bridge"}, nil },
 	})
 

diff --git a/pkg/bridge/ai/service.go b/pkg/bridge/ai/service.go
@@ -46,9 +46,9 @@ type ServiceOptions struct {
 	// CallerCallTimeout is the timeout for awaiting the function response.
 	CallerCallTimeout time.Duration
 	// SourceBuilder should builds an unconnected source.
-	SourceBuilder func() yomo.Source
+	SourceBuilder func(credential string) yomo.Source
 	// ReducerBuilder should builds an unconnected reducer.
-	ReducerBuilder func() yomo.StreamFunction
+	ReducerBuilder func(credential string) yomo.StreamFunction
 	// MetadataExchanger exchanges metadata from the credential.
 	MetadataExchanger func(credential string) (metadata.M, error)
 }
@@ -446,8 +446,8 @@ func (srv *Service) loadOrCreateCaller(credential string) (*Caller, error) {
 		return nil, err
 	}
 	caller, err = srv.newCallerFunc(
-		srv.option.SourceBuilder(),
-		srv.option.ReducerBuilder(),
+		srv.option.SourceBuilder(credential),
+		srv.option.ReducerBuilder(credential),
 		md,
 		srv.option.CallerCallTimeout,
 	)

diff --git a/pkg/bridge/ai/service_test.go b/pkg/bridge/ai/service_test.go
@@ -221,8 +221,8 @@ func TestServiceInvoke(t *testing.T) {
 			}
 
 			service := newService(pd, newCaller, &ServiceOptions{
-				SourceBuilder:     func() yomo.Source { return flow },
-				ReducerBuilder:    func() yomo.StreamFunction { return flow },
+				SourceBuilder:     func(_ string) yomo.Source { return flow },
+				ReducerBuilder:    func(_ string) yomo.StreamFunction { return flow },
 				MetadataExchanger: func(_ string) (metadata.M, error) { return metadata.M{"hello": "llm bridge"}, nil },
 			})
 
@@ -390,8 +390,8 @@ func TestServiceChatCompletion(t *testing.T) {
 			}
 
 			service := newService(pd, newCaller, &ServiceOptions{
-				SourceBuilder:     func() yomo.Source { return flow },
-				ReducerBuilder:    func() yomo.StreamFunction { return flow },
+				SourceBuilder:     func(_ string) yomo.Source { return flow },
+				ReducerBuilder:    func(_ string) yomo.StreamFunction { return flow },
 				MetadataExchanger: func(_ string) (metadata.M, error) { return metadata.M{"hello": "llm bridge"}, nil },
 			})