Repository: tmc/langchaingo
Pull request #170 (2023-07-04 23:15:37)

huggingface: build inferencePayload appropriately
### Description
When using the Hugging Face client, parameters are not being correctly included in the `inferencePayload`, causing the Inference API to ignore them. The [docs](https://huggingface.co/docs/api-inference/detailed_parameters#text-generation-task) for performing a text generation task specify that the parameters should be specified in the body of the request like so:
```json
{
    "inputs": "Can you please let us know more details about your ",
    "parameters": {
        "top_k": 10
    }
}
```

but we're currently doing this:
```json
{
    "inputs": "Can you please let us know more details about your ",
    "top_k": 10
}
```

I've been able to test locally that this proposed change works.

### PR Checklist

- [x] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [x] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [x] Name your Pull Request title clearly, concisely, and prefixed with the name of the primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [x] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [x] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `Fixes #123`).
- [x] Describes the source of new concepts.
- [x] References existing implementations as appropriate.
- [x] Contains test coverage for new functions.
- [x] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.



Patch:
diff --git a/llms/huggingface/internal/huggingfaceclient/huggingfaceclient.go b/llms/huggingface/internal/huggingfaceclient/huggingfaceclient.go
index 6a4dc9b8..53941cc5 100644
--- a/llms/huggingface/internal/huggingfaceclient/huggingfaceclient.go
+++ b/llms/huggingface/internal/huggingfaceclient/huggingfaceclient.go
@@ -14,6 +14,7 @@ var (
 type Client struct {
 	Token string
 	Model string
+	url   string
 }
 
 func New(token string, model string) (*Client, error) {
@@ -23,6 +24,7 @@ func New(token string, model string) (*Client, error) {
 	return &Client{
 		Token: token,
 		Model: model,
+		url:   hfInferenceAPI,
 	}, nil
 }
 
@@ -45,15 +47,17 @@ type InferenceResponse struct {
 
 func (c *Client) RunInference(ctx context.Context, request *InferenceRequest) (*InferenceResponse, error) {
 	payload := &inferencePayload{
-		Model:             request.Model,
-		Inputs:            request.Prompt,
-		Temperature:       request.Temperature,
-		TopP:              request.TopP,
-		TopK:              request.TopK,
-		MinLength:         request.MinLength,
-		MaxLength:         request.MaxLength,
-		RepetitionPenalty: request.RepetitionPenalty,
-		Seed:              request.Seed,
+		Model:  request.Model,
+		Inputs: request.Prompt,
+		Parameters: parameters{
+			Temperature:       request.Temperature,
+			TopP:              request.TopP,
+			TopK:              request.TopK,
+			MinLength:         request.MinLength,
+			MaxLength:         request.MaxLength,
+			RepetitionPenalty: request.RepetitionPenalty,
+			Seed:              request.Seed,
+		},
 	}
 	resp, err := c.runInference(ctx, payload)
 	if err != nil {
diff --git a/llms/huggingface/internal/huggingfaceclient/inference.go b/llms/huggingface/internal/huggingfaceclient/inference.go
index 2d9d1c3e..be4a000b 100644
--- a/llms/huggingface/internal/huggingfaceclient/inference.go
+++ b/llms/huggingface/internal/huggingfaceclient/inference.go
@@ -4,10 +4,14 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"io"
 	"net/http"
 )
 
+var ErrUnexpectedStatusCode = errors.New("unexpected status code")
+
 // InferenceTask is the type of inference task to run.
 type InferenceTask string
 
@@ -17,8 +21,12 @@ const (
 )
 
 type inferencePayload struct {
-	Model             string  `json:"-"`
-	Inputs            string  `json:"inputs"`
+	Model      string     `json:"-"`
+	Inputs     string     `json:"inputs"`
+	Parameters parameters `json:"parameters,omitempty"`
+}
+
+type parameters struct {
 	Temperature       float64 `json:"temperature,omitempty"`
 	TopP              float64 `json:"top_p,omitempty"`
 	TopK              int     `json:"top_k,omitempty"`
@@ -43,7 +51,7 @@ func (c *Client) runInference(ctx context.Context, payload *inferencePayload) (i
 		return nil, err
 	}
 	body := bytes.NewReader(payloadBytes)
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("%s%s", hfInferenceAPI, payload.Model), body) //nolint:lll
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("%s%s", c.url, payload.Model), body) //nolint:lll
 	if err != nil {
 		return nil, err
 	}
@@ -64,6 +72,20 @@ func (c *Client) runInference(ctx context.Context, payload *inferencePayload) (i
 	}
 	defer r.Body.Close()
 
+	if r.StatusCode != http.StatusOK {
+		b, err := io.ReadAll(r.Body)
+		if err != nil {
+			return nil, fmt.Errorf("failed to read response body: %w", err)
+		}
+
+		if len(b) > 0 {
+			err = fmt.Errorf("%w: %d, body: %s", ErrUnexpectedStatusCode, r.StatusCode, string(b))
+		} else {
+			err = fmt.Errorf("%w: %d", ErrUnexpectedStatusCode, r.StatusCode)
+		}
+		return nil, err
+	}
+
 	// debug print the http response with httputil:
 	// resDump, err := httputil.DumpResponse(r, true)
 	// if err != nil {


Test patch:
diff --git a/llms/huggingface/internal/huggingfaceclient/huggingfaceclient_test.go b/llms/huggingface/internal/huggingfaceclient/huggingfaceclient_test.go
new file mode 100644
index 00000000..b14e7c30
--- /dev/null
+++ b/llms/huggingface/internal/huggingfaceclient/huggingfaceclient_test.go
@@ -0,0 +1,74 @@
+package huggingfaceclient
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	goodResponse = "I hug, you hug, we all hug!"
+	errMsg       = "Error in `parameters.top_k`: ensure this value is greater than or equal to 1"
+)
+
+func TestRunInference(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		req      *InferenceRequest
+		expected *InferenceResponse
+		wantErr  string
+	}{
+		{"ok", &InferenceRequest{}, &InferenceResponse{Text: goodResponse}, ""},
+		{"not ok", &InferenceRequest{TopK: -1}, nil, errMsg},
+	}
+
+	server := mockServer(t)
+	t.Cleanup(server.Close)
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			client, err := New("token", "model")
+			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)
+			if tc.wantErr != "" {
+				assert.ErrorContains(t, err, tc.wantErr)
+			}
+		})
+	}
+}
+
+func mockServer(t *testing.T) *httptest.Server {
+	t.Helper()
+
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		b, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+
+		var infReq inferencePayload
+		err = json.Unmarshal(b, &infReq)
+		require.NoError(t, err)
+
+		if infReq.Parameters.TopK == -1 {
+			w.WriteHeader(http.StatusBadRequest)
+			_, _ = w.Write([]byte(fmt.Sprintf(`{"error":["%s"]}`, errMsg)))
+		} else {
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte(fmt.Sprintf(`[{"generated_text":"%s"}]`, goodResponse)))
+		}
+	}))
+}