65 lines
1.3 KiB
Go
65 lines
1.3 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"3588AdminBackend/internal/config"
|
|
)
|
|
|
|
type AgentClient struct {
|
|
cfg *config.Config
|
|
client *http.Client
|
|
}
|
|
|
|
func NewAgentClient(cfg *config.Config) *AgentClient {
|
|
return &AgentClient{
|
|
cfg: cfg,
|
|
client: &http.Client{
|
|
Timeout: 3 * time.Second, // Overall timeout
|
|
Transport: &http.Transport{
|
|
DialContext: (&http.Transport{}).DialContext, // Default dialer
|
|
// Connect timeout is usually handled at dialer level but 1s is tight.
|
|
// For now using the simple client timeout.
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func (c *AgentClient) Do(method, ip string, port int, path string, body []byte) ([]byte, int, error) {
|
|
url := fmt.Sprintf("http://%s:%d%s", ip, port, path)
|
|
req, err := http.NewRequest(method, url, bytes.NewBuffer(body))
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
|
|
if c.cfg.AgentToken != "" {
|
|
req.Header.Set("X-RK-Token", c.cfg.AgentToken)
|
|
}
|
|
|
|
if len(body) > 0 {
|
|
contentType := "application/json"
|
|
if strings.HasPrefix(path, "/v1/models/") {
|
|
contentType = "application/octet-stream"
|
|
}
|
|
req.Header.Set("Content-Type", contentType)
|
|
}
|
|
|
|
resp, err := c.client.Do(req)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
respBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, resp.StatusCode, err
|
|
}
|
|
|
|
return respBody, resp.StatusCode, nil
|
|
}
|