Merge pull request #545 from mheers/main

huggingface: allows to set custom inference endpoint

llms/huggingface/huggingfacellm.go

@@ -76,6 +76,7 @@ func New(opts ...Option) (*LLM, error) {
 	options := &options{
 		token: os.Getenv(tokenEnvVarName),
 		model: defaultModel,
+		url:   defaultURL,
 	}
 
 	for _, opt := range opts {
@@ -86,7 +87,7 @@ func New(opts ...Option) (*LLM, error) {
 		return nil, ErrMissingToken
 	}
 
-	c, err := huggingfaceclient.New(options.token, options.model)
+	c, err := huggingfaceclient.New(options.token, options.model, options.url)
 	if err != nil {
 		return nil, err
 	}

llms/huggingface/huggingfacellm_option.go

@@ -3,11 +3,13 @@ package huggingface
 const (
 	tokenEnvVarName = "HUGGINGFACEHUB_API_TOKEN"
 	defaultModel    = "gpt2"
+	defaultURL      = "https://api-inference.huggingface.co"
 )
 
 type options struct {
 	token string
 	model string
+	url   string
 }
 
 type Option func(*options)
@@ -27,3 +29,11 @@ func WithModel(model string) Option {
 		opts.model = model
 	}
 }
+
+// WithURL passes the HuggingFace url to the client. If not set, then will be
+// used default url.
+func WithURL(url string) Option {
+	return func(opts *options) {
+		opts.url = url
+	}
+}

llms/huggingface/internal/huggingfaceclient/huggingfaceclient.go

@@ -11,22 +11,20 @@ var (
 	ErrEmptyResponse = errors.New("empty response")
 )
 
-const huggingfaceAPIBaseURL = "https://api-inference.huggingface.co"
-
 type Client struct {
 	Token string
 	Model string
 	url   string
 }
 
-func New(token string, model string) (*Client, error) {
+func New(token, model, url string) (*Client, error) {
 	if token == "" {
 		return nil, ErrInvalidToken
 	}
 	return &Client{
 		Token: token,
 		Model: model,
-		url:   huggingfaceAPIBaseURL,
+		url:   url,
 	}, nil
 }
 

llms/huggingface/internal/huggingfaceclient/huggingfaceclient_test.go

@@ -38,10 +38,8 @@ func TestRunInference(t *testing.T) {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			client, err := New("token", "model")
+			client, err := New("token", "model", server.URL)
 			require.NoError(t, err)
-			// Override the URL to point to our mock server.
-			client.url = server.URL
 
 			resp, err := client.RunInference(context.TODO(), tc.req)
 			assert.Equal(t, tc.expected, resp)