104 lines
2.6 KiB
Go
104 lines
2.6 KiB
Go
package ai
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
type Client struct {
|
||
BaseURL string
|
||
APIKey string
|
||
Model string
|
||
Timeout time.Duration
|
||
}
|
||
|
||
type chatMessage struct {
|
||
Role string `json:"role"`
|
||
Content string `json:"content"`
|
||
}
|
||
|
||
type chatRequest struct {
|
||
Model string `json:"model"`
|
||
Messages []chatMessage `json:"messages"`
|
||
Temperature float64 `json:"temperature"`
|
||
}
|
||
|
||
type chatResponse struct {
|
||
Choices []struct {
|
||
Message chatMessage `json:"message"`
|
||
} `json:"choices"`
|
||
Error *struct {
|
||
Message string `json:"message"`
|
||
} `json:"error"`
|
||
}
|
||
|
||
func (c *Client) Suggest(userInput string) (string, error) {
|
||
return c.chat(userInput)
|
||
}
|
||
|
||
func (c *Client) Explain(result string) (string, error) {
|
||
return "", nil
|
||
}
|
||
|
||
func commandGuide() string {
|
||
b, err := os.ReadFile("docs/ai_command_guide.md")
|
||
if err != nil {
|
||
return ""
|
||
}
|
||
return strings.TrimSpace(string(b))
|
||
}
|
||
|
||
func (c *Client) chat(userInput string) (string, error) {
|
||
if strings.TrimSpace(c.BaseURL) == "" || strings.TrimSpace(c.APIKey) == "" || strings.TrimSpace(c.Model) == "" {
|
||
return "", errors.New("ai config missing")
|
||
}
|
||
base := strings.TrimRight(c.BaseURL, "/")
|
||
url := base + "/chat/completions"
|
||
|
||
sys := "你是命令翻译器。把用户的自然语言转换成系统支持的标准命令。只输出一行命令，不要解释。若无法确定，输出 FAIL。\n\n可用命令知识库：\n" + commandGuide() + "\n\n规则：严格按命令格式输出。缺少关键参数时输出 FAIL。不要猜测 zone_id/record_id/backup_id。"
|
||
req := chatRequest{
|
||
Model: c.Model,
|
||
Messages: []chatMessage{
|
||
{Role: "system", Content: sys},
|
||
{Role: "user", Content: userInput},
|
||
},
|
||
Temperature: 0,
|
||
}
|
||
body, _ := json.Marshal(req)
|
||
client := &http.Client{Timeout: c.Timeout}
|
||
httpReq, _ := http.NewRequest("POST", url, bytes.NewReader(body))
|
||
httpReq.Header.Set("Content-Type", "application/json")
|
||
httpReq.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||
|
||
resp, err := client.Do(httpReq)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
defer resp.Body.Close()
|
||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||
if resp.StatusCode == 429 {
|
||
return "", fmt.Errorf("ai rate limited")
|
||
}
|
||
return "", fmt.Errorf("ai http %d", resp.StatusCode)
|
||
}
|
||
var out chatResponse
|
||
if err := json.Unmarshal(b, &out); err != nil {
|
||
return "", err
|
||
}
|
||
if out.Error != nil && out.Error.Message != "" {
|
||
return "", errors.New(out.Error.Message)
|
||
}
|
||
if len(out.Choices) == 0 {
|
||
return "", errors.New("empty ai response")
|
||
}
|
||
return strings.TrimSpace(out.Choices[0].Message.Content), nil
|
||
}
|