feat(store): add Golang client (#1977)

This adds a basic store client for Go

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto 2024-04-16 15:54:14 +02:00 committed by GitHub
parent df4a13a08b
commit 33c78d2228
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

151
core/clients/store.go Normal file
View File

@ -0,0 +1,151 @@
package clients
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
)
// Define a struct to hold the store API client
type StoreClient struct {
BaseURL string
Client *http.Client
}
type SetRequest struct {
Keys [][]float32 `json:"keys"`
Values []string `json:"values"`
}
type GetRequest struct {
Keys [][]float32 `json:"keys"`
}
type GetResponse struct {
Keys [][]float32 `json:"keys"`
Values []string `json:"values"`
}
type DeleteRequest struct {
Keys [][]float32 `json:"keys"`
}
type FindRequest struct {
TopK int `json:"topk"`
Key []float32 `json:"key"`
}
type FindResponse struct {
Keys [][]float32 `json:"keys"`
Values []string `json:"values"`
Similarities []float32 `json:"similarities"`
}
// Constructor for StoreClient
func NewStoreClient(baseUrl string) *StoreClient {
return &StoreClient{
BaseURL: baseUrl,
Client: &http.Client{},
}
}
// Implement Set method
func (c *StoreClient) Set(req SetRequest) error {
return c.doRequest("stores/set", req)
}
// Implement Get method
func (c *StoreClient) Get(req GetRequest) (*GetResponse, error) {
body, err := c.doRequestWithResponse("stores/get", req)
if err != nil {
return nil, err
}
var resp GetResponse
err = json.Unmarshal(body, &resp)
if err != nil {
return nil, err
}
return &resp, nil
}
// Implement Delete method
func (c *StoreClient) Delete(req DeleteRequest) error {
return c.doRequest("stores/delete", req)
}
// Implement Find method
func (c *StoreClient) Find(req FindRequest) (*FindResponse, error) {
body, err := c.doRequestWithResponse("stores/find", req)
if err != nil {
return nil, err
}
var resp FindResponse
err = json.Unmarshal(body, &resp)
if err != nil {
return nil, err
}
return &resp, nil
}
// Helper function to perform a request without expecting a response body
func (c *StoreClient) doRequest(path string, data interface{}) error {
jsonData, err := json.Marshal(data)
if err != nil {
return err
}
req, err := http.NewRequest("POST", c.BaseURL+"/"+path, bytes.NewBuffer(jsonData))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.Client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("API request to %s failed with status code %d", path, resp.StatusCode)
}
return nil
}
// Helper function to perform a request and parse the response body
func (c *StoreClient) doRequestWithResponse(path string, data interface{}) ([]byte, error) {
jsonData, err := json.Marshal(data)
if err != nil {
return nil, err
}
req, err := http.NewRequest("POST", c.BaseURL+"/"+path, bytes.NewBuffer(jsonData))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.Client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request to %s failed with status code %d", path, resp.StatusCode)
}
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
return body, nil
}