156 lines
3.5 KiB
Go
156 lines
3.5 KiB
Go
package ai
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
// SiliconFlowClient 硅基流动API客户端
|
|
type SiliconFlowClient struct {
|
|
apiKey string
|
|
baseURL string
|
|
model string
|
|
timeout time.Duration
|
|
maxRetries int
|
|
}
|
|
|
|
// Config AI配置
|
|
type Config struct {
|
|
Endpoint string `yaml:"endpoint"`
|
|
Model string `yaml:"model"`
|
|
Timeout int `yaml:"timeout"`
|
|
MaxRetries int `yaml:"max_retries"`
|
|
ApiKey string `yaml:"api_key"`
|
|
}
|
|
|
|
// ChatMessage 聊天消息
|
|
type ChatMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
// ChatRequest 聊天请求
|
|
type ChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []ChatMessage `json:"messages"`
|
|
Temperature float64 `json:"temperature"`
|
|
Stream bool `json:"stream"`
|
|
}
|
|
|
|
// ChatResponse 聊天响应
|
|
type ChatResponse struct {
|
|
ID string `json:"id"`
|
|
Object string `json:"object"`
|
|
Created int64 `json:"created"`
|
|
Model string `json:"model"`
|
|
Choices []Choice `json:"choices"`
|
|
Usage Usage `json:"usage"`
|
|
}
|
|
|
|
// Choice 选择项
|
|
type Choice struct {
|
|
Index int `json:"index"`
|
|
Message ChatMessage `json:"message"`
|
|
FinishReason string `json:"finish_reason"`
|
|
}
|
|
|
|
// Usage 使用情况
|
|
type Usage struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
CompletionTokens int `json:"completion_tokens"`
|
|
TotalTokens int `json:"total_tokens"`
|
|
}
|
|
|
|
// NewClient 创建新的AI客户端
|
|
func NewClient(config Config) *SiliconFlowClient {
|
|
return &SiliconFlowClient{
|
|
apiKey: config.ApiKey,
|
|
baseURL: config.Endpoint,
|
|
model: config.Model,
|
|
timeout: time.Duration(config.Timeout) * time.Second,
|
|
maxRetries: config.MaxRetries,
|
|
}
|
|
}
|
|
|
|
// Generate 生成内容
|
|
func (c *SiliconFlowClient) Generate(prompt string) (string, error) {
|
|
reqBody := ChatRequest{
|
|
Model: c.model,
|
|
Messages: []ChatMessage{
|
|
{
|
|
Role: "user",
|
|
Content: prompt,
|
|
},
|
|
},
|
|
Temperature: 0.7,
|
|
Stream: false,
|
|
}
|
|
|
|
reqJSON, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to marshal request: %w", err)
|
|
}
|
|
|
|
var lastErr error
|
|
for attempt := 0; attempt < c.maxRetries; attempt++ {
|
|
result, err := c.sendRequest(reqJSON)
|
|
if err == nil {
|
|
return result, nil
|
|
}
|
|
|
|
lastErr = err
|
|
|
|
if attempt < c.maxRetries-1 {
|
|
time.Sleep(time.Duration(attempt+1) * time.Second)
|
|
}
|
|
}
|
|
|
|
return "", fmt.Errorf("failed after %d attempts: %w", c.maxRetries, lastErr)
|
|
}
|
|
|
|
func (c *SiliconFlowClient) sendRequest(reqJSON []byte) (string, error) {
|
|
url := c.baseURL + "/chat/completions"
|
|
|
|
req, err := http.NewRequest("POST", url, bytes.NewReader(reqJSON))
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create request: %w", err)
|
|
}
|
|
|
|
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
client := &http.Client{
|
|
Timeout: c.timeout,
|
|
}
|
|
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to send request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to read response: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return "", fmt.Errorf("API returned status code %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
var chatResp ChatResponse
|
|
if err := json.Unmarshal(body, &chatResp); err != nil {
|
|
return "", fmt.Errorf("failed to unmarshal response: %w", err)
|
|
}
|
|
|
|
if len(chatResp.Choices) == 0 {
|
|
return "", fmt.Errorf("no choices in response")
|
|
}
|
|
|
|
return chatResp.Choices[0].Message.Content, nil
|
|
}
|