118 lines
2.9 KiB
Go
118 lines
2.9 KiB
Go
package ai
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// openAICompatible 是 OpenAI Chat Completions 兼容端点的通用流式客户端；
|
|
// 讯飞星火 Lite 与 DeepSeek 均暴露该协议，只是 baseURL/model/key 不同。
|
|
type openAICompatible struct {
|
|
name string
|
|
baseURL string
|
|
model string
|
|
apiKey string
|
|
client *http.Client
|
|
}
|
|
|
|
func (p *openAICompatible) Name() string { return p.name }
|
|
|
|
type chatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []Message `json:"messages"`
|
|
Stream bool `json:"stream"`
|
|
}
|
|
|
|
type chatDelta struct {
|
|
Choices []struct {
|
|
Delta struct {
|
|
Content string `json:"content"`
|
|
} `json:"delta"`
|
|
} `json:"choices"`
|
|
Error *struct {
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
}
|
|
|
|
func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (<-chan Chunk, error) {
|
|
body, e := json.Marshal(chatRequest{Model: p.model, Messages: messages, Stream: true})
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
req, e := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(p.baseURL, "/")+"/chat/completions", bytes.NewReader(body))
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
|
req.Header.Set("Accept", "text/event-stream")
|
|
resp, e := p.client.Do(req)
|
|
if e != nil {
|
|
return nil, errors.New("AI_NETWORK_ERROR")
|
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
|
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
|
resp.Body.Close()
|
|
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
|
|
return nil, errors.New("AI_KEY_INVALID")
|
|
}
|
|
return nil, fmt.Errorf("AI_HTTP_%d: %s", resp.StatusCode, truncate(string(raw), 300))
|
|
}
|
|
|
|
out := make(chan Chunk, 16)
|
|
go func() {
|
|
defer close(out)
|
|
defer resp.Body.Close()
|
|
sc := bufio.NewScanner(resp.Body)
|
|
sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
|
for sc.Scan() {
|
|
line := strings.TrimSpace(sc.Text())
|
|
if !strings.HasPrefix(line, "data:") {
|
|
continue
|
|
}
|
|
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
|
if payload == "" || payload == "[DONE]" {
|
|
continue
|
|
}
|
|
var d chatDelta
|
|
if json.Unmarshal([]byte(payload), &d) != nil {
|
|
continue
|
|
}
|
|
if d.Error != nil {
|
|
out <- Chunk{Err: errors.New(truncate(d.Error.Message, 300))}
|
|
return
|
|
}
|
|
if len(d.Choices) > 0 && d.Choices[0].Delta.Content != "" {
|
|
select {
|
|
case out <- Chunk{Content: d.Choices[0].Delta.Content}:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
if e := sc.Err(); e != nil && ctx.Err() == nil {
|
|
out <- Chunk{Err: errors.New("AI_STREAM_INTERRUPTED")}
|
|
}
|
|
}()
|
|
return out, nil
|
|
}
|
|
|
|
func truncate(s string, n int) string {
|
|
if len(s) <= n {
|
|
return s
|
|
}
|
|
return s[:n] + "..."
|
|
}
|
|
|
|
func newHTTPClient() *http.Client {
|
|
return &http.Client{Timeout: 5 * time.Minute}
|
|
}
|