139 lines
3.9 KiB
Go
139 lines
3.9 KiB
Go
// OpenAI 兼容出站适配器：IR → OpenAI chat/completions（元炁亦走此协议，type=local 复用）。
|
||
package transform
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
|
||
"hunjing/shenshi/internal/ir"
|
||
)
|
||
|
||
// UpstreamError 上游错误：携带 HTTP 状态，供路由层分类冷却（凭证级/绑定级）。
|
||
type UpstreamError struct {
|
||
StatusCode int
|
||
Message string
|
||
}
|
||
|
||
func (e *UpstreamError) Error() string {
|
||
return fmt.Sprintf("上游 %d: %s", e.StatusCode, e.Message)
|
||
}
|
||
|
||
// OpenAICompatAdaptor OpenAI 兼容端点适配器（云厂商兼容面 + 元炁本地端点）。
|
||
type OpenAICompatAdaptor struct {
|
||
EndpointName string // 端点标识（usage 日志与冷却键）
|
||
BaseURL string
|
||
APIKey string
|
||
// UpstreamModel 该端点承载的上游模型名（逻辑模型 → 端点绑定的组成部分）
|
||
UpstreamModel string
|
||
Client *http.Client
|
||
}
|
||
|
||
// Name 实现 Adaptor。
|
||
func (a *OpenAICompatAdaptor) Name() string { return a.EndpointName }
|
||
|
||
type oaMessage struct {
|
||
Role string `json:"role"`
|
||
Content string `json:"content"`
|
||
}
|
||
|
||
type oaRequest struct {
|
||
Model string `json:"model"`
|
||
Messages []oaMessage `json:"messages"`
|
||
Stream bool `json:"stream"`
|
||
Temperature *float64 `json:"temperature,omitempty"`
|
||
MaxTokens *int `json:"max_tokens,omitempty"`
|
||
}
|
||
|
||
type oaChunk struct {
|
||
Choices []struct {
|
||
Delta struct {
|
||
Content string `json:"content"`
|
||
} `json:"delta"`
|
||
FinishReason *string `json:"finish_reason"`
|
||
} `json:"choices"`
|
||
Usage *struct {
|
||
PromptTokens int `json:"prompt_tokens"`
|
||
CompletionTokens int `json:"completion_tokens"`
|
||
} `json:"usage"`
|
||
}
|
||
|
||
// Stream 实现 Adaptor：SSE 逐帧转 IR Delta。首字节超时由调用方 context 控制。
|
||
func (a *OpenAICompatAdaptor) Stream(ctx context.Context, req *ir.ChatRequest, out chan<- ir.Delta) error {
|
||
body := oaRequest{
|
||
Model: a.UpstreamModel,
|
||
Messages: make([]oaMessage, 0, len(req.Messages)),
|
||
Stream: true,
|
||
}
|
||
for _, m := range req.Messages {
|
||
body.Messages = append(body.Messages, oaMessage{Role: m.Role, Content: m.Content})
|
||
}
|
||
body.Temperature = req.Temperature
|
||
body.MaxTokens = req.MaxTokens
|
||
|
||
payload, err := json.Marshal(body)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(a.BaseURL, "/")+"/v1/chat/completions", bytes.NewReader(payload))
|
||
if err != nil {
|
||
return err
|
||
}
|
||
httpReq.Header.Set("content-type", "application/json")
|
||
if a.APIKey != "" {
|
||
httpReq.Header.Set("authorization", "Bearer "+a.APIKey)
|
||
}
|
||
|
||
client := a.Client
|
||
if client == nil {
|
||
client = &http.Client{Timeout: 120 * time.Second}
|
||
}
|
||
resp, err := client.Do(httpReq)
|
||
if err != nil {
|
||
return &UpstreamError{StatusCode: 0, Message: err.Error()} // 0 = 网络/超时（绑定级冷却）
|
||
}
|
||
defer func() { _ = resp.Body.Close() }()
|
||
if resp.StatusCode >= 400 {
|
||
buf := new(bytes.Buffer)
|
||
_, _ = buf.ReadFrom(resp.Body)
|
||
return &UpstreamError{StatusCode: resp.StatusCode, Message: buf.String()}
|
||
}
|
||
|
||
scanner := bufio.NewScanner(resp.Body)
|
||
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
||
for scanner.Scan() {
|
||
line := scanner.Text()
|
||
if !strings.HasPrefix(line, "data:") {
|
||
continue
|
||
}
|
||
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||
if data == "" || data == "[DONE]" {
|
||
continue
|
||
}
|
||
var chunk oaChunk
|
||
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
||
continue
|
||
}
|
||
if len(chunk.Choices) > 0 {
|
||
c := chunk.Choices[0]
|
||
if c.Delta.Content != "" {
|
||
out <- ir.Delta{Delta: c.Delta.Content}
|
||
}
|
||
if c.FinishReason != nil {
|
||
finish := *c.FinishReason
|
||
usagePrompt, usageCompletion := 0, 0
|
||
if chunk.Usage != nil {
|
||
usagePrompt, usageCompletion = chunk.Usage.PromptTokens, chunk.Usage.CompletionTokens
|
||
}
|
||
out <- ir.Delta{FinishReason: finish, UsagePromptTokens: usagePrompt, UsageCompletionTokens: usageCompletion}
|
||
}
|
||
}
|
||
}
|
||
return scanner.Err()
|
||
}
|