103 lines
2.1 KiB
Go
103 lines
2.1 KiB
Go
package client
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
)
|
|
|
|
type EmbeddingsClient struct {
|
|
url string
|
|
apiKey string
|
|
|
|
Do func(req *http.Request) (*http.Response, error)
|
|
}
|
|
|
|
type EmbeddingsRequest struct {
|
|
Model string `json:"model"`
|
|
Input string `json:"input"`
|
|
}
|
|
|
|
type EmbeddingsData struct {
|
|
Object string `json:"object"`
|
|
Index int `json:"index"`
|
|
Embeddings []float64 `json:"embedding"`
|
|
}
|
|
|
|
type EmbeddingsResponse struct {
|
|
Object string `json:"object"`
|
|
Data []EmbeddingsData `json:"data"`
|
|
}
|
|
|
|
func CreateEmbeddingsClient() EmbeddingsClient {
|
|
apiKey := os.Getenv(OPENAI_API_KEY)
|
|
|
|
if len(apiKey) == 0 {
|
|
panic("No api key")
|
|
}
|
|
|
|
return EmbeddingsClient{
|
|
apiKey: apiKey,
|
|
url: "https://api.openai.com/v1/embeddings",
|
|
Do: func(req *http.Request) (*http.Response, error) {
|
|
client := &http.Client{}
|
|
return client.Do(req)
|
|
},
|
|
}
|
|
}
|
|
|
|
func (client EmbeddingsClient) getRequest(text string) (*http.Request, error) {
|
|
embeddingsReq := EmbeddingsRequest{
|
|
Model: "text-embedding-3-large",
|
|
Input: text,
|
|
}
|
|
|
|
jsonEmbeddingsBody, err := json.Marshal(embeddingsReq)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req, err := http.NewRequest("POST", client.url, bytes.NewBuffer(jsonEmbeddingsBody))
|
|
if err != nil {
|
|
return req, err
|
|
}
|
|
|
|
req.Header.Add("Authorization", "Bearer "+client.apiKey)
|
|
req.Header.Add("Content-Type", "application/json")
|
|
|
|
return req, nil
|
|
}
|
|
|
|
func (client EmbeddingsClient) Request(text string) (EmbeddingsData, error) {
|
|
httpRequest, err := client.getRequest(text)
|
|
if err != nil {
|
|
return EmbeddingsData{}, err
|
|
}
|
|
|
|
resp, err := client.Do(httpRequest)
|
|
if err != nil {
|
|
return EmbeddingsData{}, err
|
|
}
|
|
|
|
response, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return EmbeddingsData{}, err
|
|
}
|
|
|
|
embeddingsResponse := EmbeddingsResponse{}
|
|
err = json.Unmarshal(response, &embeddingsResponse)
|
|
|
|
if err != nil {
|
|
return EmbeddingsData{}, err
|
|
}
|
|
|
|
if len(embeddingsResponse.Data) != 1 {
|
|
return EmbeddingsData{}, errors.New("Unsupported. We currently only accept 1 choice from AI.")
|
|
}
|
|
|
|
return embeddingsResponse.Data[0], nil
|
|
}
|