Spaces:
Running
Running
File size: 5,218 Bytes
b110593 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
// _ _
// __ _____ __ ___ ___ __ _| |_ ___
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
// \ V V / __/ (_| |\ V /| | (_| | || __/
// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
//
// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
//
// CONTACT: [email protected]
//
package clients
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/weaviate/weaviate/usecases/modulecomponents"
"github.com/pkg/errors"
"github.com/sirupsen/logrus"
"github.com/weaviate/weaviate/modules/text2vec-cohere/ent"
)
type embeddingsRequest struct {
Texts []string `json:"texts"`
Model string `json:"model,omitempty"`
Truncate string `json:"truncate,omitempty"`
InputType inputType `json:"input_type,omitempty"`
}
type embeddingsResponse struct {
Embeddings [][]float32 `json:"embeddings,omitempty"`
Message string `json:"message,omitempty"`
}
type vectorizer struct {
apiKey string
httpClient *http.Client
urlBuilder *cohereUrlBuilder
logger logrus.FieldLogger
}
type inputType string
const (
searchDocument inputType = "search_document"
searchQuery inputType = "search_query"
)
func New(apiKey string, timeout time.Duration, logger logrus.FieldLogger) *vectorizer {
return &vectorizer{
apiKey: apiKey,
httpClient: &http.Client{
Timeout: timeout,
},
urlBuilder: newCohereUrlBuilder(),
logger: logger,
}
}
func (v *vectorizer) Vectorize(ctx context.Context, input []string,
config ent.VectorizationConfig,
) (*ent.VectorizationResult, error) {
return v.vectorize(ctx, input, config.Model, config.Truncate, config.BaseURL, searchDocument)
}
func (v *vectorizer) VectorizeQuery(ctx context.Context, input []string,
config ent.VectorizationConfig,
) (*ent.VectorizationResult, error) {
return v.vectorize(ctx, input, config.Model, config.Truncate, config.BaseURL, searchQuery)
}
func (v *vectorizer) vectorize(ctx context.Context, input []string,
model, truncate, baseURL string, inputType inputType,
) (*ent.VectorizationResult, error) {
body, err := json.Marshal(embeddingsRequest{
Texts: input,
Model: model,
Truncate: truncate,
InputType: inputType,
})
if err != nil {
return nil, errors.Wrapf(err, "marshal body")
}
url := v.getCohereUrl(ctx, baseURL)
req, err := http.NewRequestWithContext(ctx, "POST", url,
bytes.NewReader(body))
if err != nil {
return nil, errors.Wrap(err, "create POST request")
}
apiKey, err := v.getApiKey(ctx)
if err != nil {
return nil, errors.Wrapf(err, "Cohere API Key")
}
req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", apiKey))
req.Header.Add("Content-Type", "application/json")
res, err := v.httpClient.Do(req)
if err != nil {
return nil, errors.Wrap(err, "send POST request")
}
defer res.Body.Close()
bodyBytes, err := io.ReadAll(res.Body)
if err != nil {
return nil, errors.Wrap(err, "read response body")
}
var resBody embeddingsResponse
if err := json.Unmarshal(bodyBytes, &resBody); err != nil {
return nil, errors.Wrap(err, "unmarshal response body")
}
if res.StatusCode >= 500 {
errorMessage := getErrorMessage(res.StatusCode, resBody.Message, "connection to Cohere failed with status: %d error: %v")
return nil, errors.Errorf(errorMessage)
} else if res.StatusCode > 200 {
errorMessage := getErrorMessage(res.StatusCode, resBody.Message, "failed with status: %d error: %v")
return nil, errors.Errorf(errorMessage)
}
if len(resBody.Embeddings) == 0 {
return nil, errors.Errorf("empty embeddings response")
}
return &ent.VectorizationResult{
Text: input,
Dimensions: len(resBody.Embeddings[0]),
Vectors: resBody.Embeddings,
}, nil
}
func (v *vectorizer) getCohereUrl(ctx context.Context, baseURL string) string {
passedBaseURL := baseURL
if headerBaseURL := v.getValueFromContext(ctx, "X-Cohere-Baseurl"); headerBaseURL != "" {
passedBaseURL = headerBaseURL
}
return v.urlBuilder.url(passedBaseURL)
}
func getErrorMessage(statusCode int, resBodyError string, errorTemplate string) string {
if resBodyError != "" {
return fmt.Sprintf(errorTemplate, statusCode, resBodyError)
}
return fmt.Sprintf(errorTemplate, statusCode)
}
func (v *vectorizer) getValueFromContext(ctx context.Context, key string) string {
if value := ctx.Value(key); value != nil {
if keyHeader, ok := value.([]string); ok && len(keyHeader) > 0 && len(keyHeader[0]) > 0 {
return keyHeader[0]
}
}
// try getting header from GRPC if not successful
if apiKey := modulecomponents.GetValueFromGRPC(ctx, key); len(apiKey) > 0 && len(apiKey[0]) > 0 {
return apiKey[0]
}
return ""
}
func (v *vectorizer) getApiKey(ctx context.Context) (string, error) {
if apiKey := v.getValueFromContext(ctx, "X-Cohere-Api-Key"); apiKey != "" {
return apiKey, nil
}
if v.apiKey != "" {
return v.apiKey, nil
}
return "", errors.New("no api key found " +
"neither in request header: X-Cohere-Api-Key " +
"nor in environment variable under COHERE_APIKEY")
}
|