98 lines
2.4 KiB
Go
98 lines
2.4 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"time"
|
||
)
|
||
|
||
// omlx 调本机 oMLX（OpenAI 兼容 /v1/chat/completions）。
|
||
// 注意：Qwen3.5 思考链默认开启且思考文本混入 content，必须传 chat_template_kwargs.enable_thinking=false。
|
||
type omlx struct {
|
||
endpoint string
|
||
apiKey string
|
||
model string
|
||
timeout int // 秒
|
||
maxTokens int
|
||
retries int
|
||
}
|
||
|
||
type chatRequest struct {
|
||
Model string `json:"model"`
|
||
Messages []chatMessage `json:"messages"`
|
||
Stream bool `json:"stream"`
|
||
MaxTokens int `json:"max_tokens"`
|
||
ChatTemplateKwargs map[string]any `json:"chat_template_kwargs"`
|
||
}
|
||
|
||
type chatMessage struct {
|
||
Role string `json:"role"`
|
||
Content string `json:"content"`
|
||
}
|
||
|
||
type chatResponse struct {
|
||
Choices []struct {
|
||
Message struct {
|
||
Content string `json:"content"`
|
||
} `json:"message"`
|
||
} `json:"choices"`
|
||
}
|
||
|
||
func (m *omlx) chat(ctx context.Context, system, user string) (string, error) {
|
||
payload := chatRequest{
|
||
Model: m.model,
|
||
Messages: []chatMessage{
|
||
{Role: "system", Content: system},
|
||
{Role: "user", Content: user},
|
||
},
|
||
Stream: false,
|
||
MaxTokens: m.maxTokens,
|
||
ChatTemplateKwargs: map[string]any{"enable_thinking": false},
|
||
}
|
||
buf, err := json.Marshal(payload)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
var lastErr error
|
||
for i := 0; i <= m.retries; i++ {
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, m.endpoint+"/v1/chat/completions", bytes.NewReader(buf))
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("Authorization", "Bearer "+m.apiKey)
|
||
resp, err := (&http.Client{Timeout: time.Duration(m.timeout) * time.Second}).Do(req)
|
||
if err != nil {
|
||
lastErr = fmt.Errorf("oMLX 请求失败: %w", err)
|
||
time.Sleep(2 * time.Second)
|
||
continue
|
||
}
|
||
body, err := io.ReadAll(resp.Body)
|
||
resp.Body.Close()
|
||
if err != nil {
|
||
lastErr = err
|
||
continue
|
||
}
|
||
if resp.StatusCode >= 400 {
|
||
lastErr = fmt.Errorf("oMLX %d: %s", resp.StatusCode, string(body))
|
||
time.Sleep(2 * time.Second)
|
||
continue
|
||
}
|
||
var cr chatResponse
|
||
if err := json.Unmarshal(body, &cr); err != nil {
|
||
lastErr = fmt.Errorf("oMLX 响应解析失败: %w", err)
|
||
continue
|
||
}
|
||
if len(cr.Choices) == 0 {
|
||
lastErr = fmt.Errorf("oMLX 空响应")
|
||
continue
|
||
}
|
||
return cr.Choices[0].Message.Content, nil
|
||
}
|
||
return "", lastErr
|
||
}
|