weaviate
99 строк · 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 clients
13
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/text2vec-gpt4all/ent"
26)
27
28type client struct {
29origin string
30httpClient *http.Client
31logger logrus.FieldLogger
32}
33
34func New(origin string, timeout time.Duration, logger logrus.FieldLogger) *client {
35return &client{
36origin: origin,
37httpClient: &http.Client{
38Timeout: timeout,
39},
40logger: logger,
41}
42}
43
44func (c *client) Vectorize(ctx context.Context, text string) (*ent.VectorizationResult, error) {
45body, err := json.Marshal(vecRequest{text})
46if err != nil {
47return nil, errors.Wrapf(err, "marshal body")
48}
49
50req, err := http.NewRequestWithContext(ctx, "POST", c.url("/vectorize"),
51bytes.NewReader(body))
52if err != nil {
53return nil, errors.Wrap(err, "create POST request")
54}
55
56res, err := c.httpClient.Do(req)
57if err != nil {
58return nil, errors.Wrap(err, "send POST request")
59}
60defer res.Body.Close()
61
62bodyBytes, err := io.ReadAll(res.Body)
63if err != nil {
64return nil, errors.Wrap(err, "read response body")
65}
66
67var resBody vecResponse
68if err := json.Unmarshal(bodyBytes, &resBody); err != nil {
69return nil, errors.Wrap(err, "unmarshal response body")
70}
71
72if res.StatusCode != 200 {
73if resBody.Error != "" {
74return nil, errors.Errorf("failed with status: %d error: %v", res.StatusCode, resBody.Error)
75}
76return nil, errors.Errorf("failed with status: %d", res.StatusCode)
77}
78
79return &ent.VectorizationResult{
80Vector: resBody.Vector,
81Dimensions: resBody.Dim,
82Text: resBody.Text,
83}, nil
84}
85
86func (c *client) url(path string) string {
87return fmt.Sprintf("%s%s", c.origin, path)
88}
89
90type vecRequest struct {
91Text string `json:"text"`
92}
93
94type vecResponse struct {
95Text string `json:"text"`
96Vector []float32 `json:"vector"`
97Dim int `json:"dim"`
98Error string `json:"error"`
99}
100