74 lines
2.7 KiB
Go
74 lines
2.7 KiB
Go
package ai
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
|
|
"write-server/internal/config"
|
|
)
|
|
|
|
var AI_PROMPTS = map[string]string{
|
|
"polish": "\u4f60\u662f\u4e00\u4e2a\u7f16\u8f91\u3002\u6da6\u8272\u4ee5\u4e0b\u6587\u672c\uff0c\u4f7f\u5176\u66f4\u6d41\u7545\u3001\u66f4\u4e13\u4e1a\u3002\u4fdd\u6301\u539f\u610f\u548c\u8bed\u6c14\uff0c\u53ea\u8fd4\u56de\u6da6\u8272\u540e\u7684\u6587\u672c\uff0c\u4e0d\u8981\u89e3\u91ca\u3002",
|
|
"expand": "\u4f60\u662f\u4e00\u4e2a\u7f16\u8f91\u3002\u6269\u5c55\u4ee5\u4e0b\u6587\u672c\uff0c\u6dfb\u52a0\u66f4\u591a\u7ec6\u8282\u548c\u89e3\u91ca\uff0c\u540c\u65f6\u4fdd\u6301\u6e05\u6670\u548c\u5f15\u4eba\u5165\u80dc\u3002\u53ea\u8fd4\u56de\u6269\u5c55\u540e\u7684\u6587\u672c\uff0c\u4e0d\u8981\u89e3\u91ca\u3002",
|
|
"shorten": "\u4f60\u662f\u4e00\u4e2a\u7f16\u8f91\u3002\u7f29\u77ed\u4ee5\u4e0b\u6587\u672c\uff0c\u4fdd\u6301\u6838\u5fc3\u610f\u601d\uff0c\u4f7f\u8868\u8fbe\u66f4\u7b80\u6d01\u3002\u53ea\u8fd4\u56de\u7f29\u77ed\u540e\u7684\u6587\u672c\uff0c\u4e0d\u8981\u89e3\u91ca\u3002",
|
|
"fix": "\u4f60\u662f\u4e00\u4e2a\u7f16\u8f91\u3002\u4fee\u590d\u4ee5\u4e0b\u6587\u672c\u4e2d\u7684\u8bed\u6cd5\u9519\u8bef\u3001\u62fc\u5199\u9519\u8bef\u548c\u4e0d\u81ea\u7136\u8868\u8fbe\uff0c\u4fdd\u6301\u539f\u610f\u3002\u53ea\u8fd4\u56de\u4fee\u6b63\u540e\u7684\u6587\u672c\uff0c\u4e0d\u8981\u89e3\u91ca\u3002",
|
|
}
|
|
|
|
type ChatMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
type ChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []ChatMessage `json:"messages"`
|
|
}
|
|
|
|
type ChatResponse struct {
|
|
Choices []struct {
|
|
Message ChatMessage `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
func AIChat(systemPrompt, userMessage string) (string, error) {
|
|
if config.DEEPSEEK_API_KEY == "" {
|
|
return "", fmt.Errorf("DEEPSEEK_API_KEY not configured")
|
|
}
|
|
|
|
req := ChatRequest{
|
|
Model: config.DEEPSEEK_MODEL,
|
|
Messages: []ChatMessage{
|
|
{Role: "system", Content: systemPrompt},
|
|
{Role: "user", Content: userMessage},
|
|
},
|
|
}
|
|
|
|
data, _ := json.Marshal(req)
|
|
httpReq, _ := http.NewRequest("POST", config.DEEPSEEK_BASE_URL+"/chat/completions", bytes.NewReader(data))
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
httpReq.Header.Set("Authorization", "Bearer "+config.DEEPSEEK_API_KEY)
|
|
|
|
resp, err := http.DefaultClient.Do(httpReq)
|
|
if err != nil {
|
|
return "", fmt.Errorf("AI request failed: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != 200 {
|
|
return "", fmt.Errorf("AI error: %s", string(body))
|
|
}
|
|
|
|
var chatResp ChatResponse
|
|
if err := json.Unmarshal(body, &chatResp); err != nil {
|
|
return "", err
|
|
}
|
|
if len(chatResp.Choices) == 0 {
|
|
return "", fmt.Errorf("no response from AI")
|
|
}
|
|
return chatResp.Choices[0].Message.Content, nil
|
|
}
|