weaviate
103 строки · 2.2 Кб
1// _ _
2// __ _____ __ ___ ___ __ _| |_ ___
3// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
4// \ V V / __/ (_| |\ V /| | (_| | || __/
5// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
6//
7// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
8//
9// CONTACT: hello@weaviate.io
10//
11
12package clients13
14import (15"bytes"16"context"17"encoding/json"18"fmt"19"io"20"net/http"21"time"22
23"github.com/pkg/errors"24"github.com/sirupsen/logrus"25"github.com/weaviate/weaviate/modules/img2vec-neural/ent"26)
27
28type vectorizer struct {29origin string30httpClient *http.Client31logger logrus.FieldLogger32}
33
34func New(origin string, timeout time.Duration, logger logrus.FieldLogger) *vectorizer {35return &vectorizer{36origin: origin,37httpClient: &http.Client{38Timeout: timeout,39},40logger: logger,41}42}
43
44func (v *vectorizer) Vectorize(ctx context.Context,45id, image string,46) (*ent.VectorizationResult, error) {47body, err := json.Marshal(vecRequest{48ID: id,49Image: image,50})51if err != nil {52return nil, errors.Wrapf(err, "marshal body")53}54
55req, err := http.NewRequestWithContext(ctx, "POST", v.url("/vectors"),56bytes.NewReader(body))57if err != nil {58return nil, errors.Wrap(err, "create POST request")59}60
61res, err := v.httpClient.Do(req)62if err != nil {63return nil, errors.Wrap(err, "send POST request")64}65defer res.Body.Close()66
67bodyBytes, err := io.ReadAll(res.Body)68if err != nil {69return nil, errors.Wrap(err, "read response body")70}71
72var resBody vecResponse73if err := json.Unmarshal(bodyBytes, &resBody); err != nil {74return nil, errors.Wrap(err, "unmarshal response body")75}76
77if res.StatusCode > 399 {78return nil, errors.Errorf("fail with status %d", res.StatusCode)79}80
81return &ent.VectorizationResult{82ID: resBody.ID,83Image: image,84Dimensions: resBody.Dim,85Vector: resBody.Vector,86}, nil87}
88
89func (v *vectorizer) url(path string) string {90return fmt.Sprintf("%s%s", v.origin, path)91}
92
93type vecRequest struct {94ID string `json:"id"`95Image string `json:"image"`96}
97
98type vecResponse struct {99ID string `json:"id"`100Vector []float32 `json:"vector"`101Dim int `json:"dim"`102Error string `json:"error"`103}
104