helper function

pluralsh · Feb 24, 2025 · e9a7b4b · e9a7b4b
1 parent fb8f403
commit e9a7b4b
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/go/ai-proxy/args/args.go b/go/ai-proxy/args/args.go
@@ -112,3 +112,7 @@ func Address() string {
 
 	return fmt.Sprintf("%s:%d", *argAddress, *argPort)
 }
+
+func OpenAICompatible() bool {
+	return Provider() == api.ProviderOpenAI || Provider() == api.ProviderBedrock
+}
diff --git a/go/ai-proxy/main.go b/go/ai-proxy/main.go
@@ -32,7 +32,7 @@ func main() {
 		router.HandleFunc(ollama.EndpointChat, p.Proxy())
 	}
 
-	if args.Provider() == api.ProviderOpenAI || args.Provider() == api.ProviderBedrock {
+	if args.OpenAICompatible() {
 		op, err := proxy.NewOpenAIProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
 		if err != nil {
 			klog.ErrorS(err, "Could not create proxy")

diff --git a/go/ai-proxy/test/helpers/common.go b/go/ai-proxy/test/helpers/common.go
@@ -35,7 +35,7 @@ func SetupServer() (*httptest.Server, error) {
 		router.HandleFunc(ollama.EndpointChat, p.Proxy())
 	}
 
-	if args.Provider() == api.ProviderOpenAI || args.Provider() == api.ProviderBedrock {
+	if args.OpenAICompatible() {
 		op, err := proxy.NewOpenAIProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
 		if err != nil {
 			klog.ErrorS(err, "Could not create proxy")