mirror of
https://github.com/mudler/LocalAI.git
synced 2024-12-18 20:27:57 +00:00
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:
parent
df4a13a08b
commit
33c78d2228
151
core/clients/store.go
Normal file
151
core/clients/store.go
Normal 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
|
||||||
|
}
|
Loading…
Reference in New Issue
Block a user