M4a:神识全量网关能力——OpenAI 兼容适配器(元炁共用)、候选路由+凭证/绑定级冷却+首字超时切道+粘性会话、用量计量与 /me/usage,故障注入测试四场景
This commit is contained in:
1 parent
df5caa69e7
commit
7279904043
8 files changed
+716
-34
No files matched your search
@@ -680,6 +680,11 @@ SHENSHI_SERVICE_TOKEN=
|
|||||||
SHENSHI_TIMEOUT_SECONDS=120.0
|
SHENSHI_TIMEOUT_SECONDS=120.0
|
||||||
# 神识服务监听端口(shenshi 服务自身)
|
# 神识服务监听端口(shenshi 服务自身)
|
||||||
SHENSHI_PORT=8080
|
SHENSHI_PORT=8080
|
||||||
|
# 神识出站端点目录(JSON 数组;未配置时回声桩兜底)。
|
||||||
|
# type: openai_compatible(云厂商兼容面与元炁 http://hunjing-yuanqi:11434 共用)
|
||||||
|
# 示例:
|
||||||
|
# SHENSHI_ENDPOINTS=[{"type":"openai_compatible","name":"yuanqi","base_url":"http://hunjing-yuanqi:11434","model":"bge-chat","priority":100}]
|
||||||
|
SHENSHI_ENDPOINTS=
|
||||||
|
|
||||||
|
|
||||||
# =====================================================================
|
# =====================================================================
|
||||||
|
|||||||
+7
-5
@@ -2,12 +2,14 @@
|
|||||||
|
|
||||||
> 文化命名「神识」(契约词典 §2.1):统一模型能力——Provider/Endpoint 接入、能力槽位、协议转换、路由、流式、BYOK、配额与成本。神识不拥有人物状态。
|
> 文化命名「神识」(契约词典 §2.1):统一模型能力——Provider/Endpoint 接入、能力槽位、协议转换、路由、流式、BYOK、配额与成本。神识不拥有人物状态。
|
||||||
|
|
||||||
## 当前状态(rewrite/1.0 M0 骨架)
|
## 当前状态(rewrite/1.0 M4a)
|
||||||
|
|
||||||
- 数据面 `POST /v1/chat/completions`:接受灵台内部 IR(`internal/ir`),经路由选适配器,以 SSE 流出标准化增量(`data: <Delta JSON>` + `data: [DONE]`);
|
- 数据面 `POST /v1/chat/completions`:IR 入站 → 候选路由 → SSE 出站;**首帧到达前不提交响应**,失败可整体切道并返回 JSON 错误(HJ-2001),已写出字节后不换道、如实 error 终帧;
|
||||||
- 当前唯一适配器为 `stub` 回声桩(`internal/transform/stub.go`),供端到端骨架验证;云 Provider(openai_compatible/anthropic)与本地元炁端点按 [重写执行计划](../docs/operations/rewrite-plan.md) M4 落地;
|
- 适配器:`openai_compatible`(云厂商兼容面与元炁本地端点共用此协议)与 `stub` 回声桩(未配置端点时兜底);
|
||||||
- 健康面 `GET /healthz` `GET /readyz`;
|
- 端点配置:环境 `SHENSHI_ENDPOINTS`(JSON 数组,字段 type/name/base_url/api_key/model/priority),示例见 `.env.example`;
|
||||||
- 服务令牌:`SHENSHI_SERVICE_TOKEN` 配置时校验 `Authorization: Bearer`(开发档可空)。
|
- 故障策略(shenshi-gateway.md §4):候选按 Priority 升序 + 粘性会话置首 + 冷却端点排尾兜底;**首字超时(默认 10s)切道**;401/403 凭证级冷却 10 分钟,429/5xx/超时仅冷却该端点绑定 30 秒;
|
||||||
|
- 计量(M4a):每次尝试一条 stdout JSON 日志(`kind=shenshi_usage`)+ 进程内聚合,`GET /me/usage` 返回按 模型×端点 的请求数/成功数/token;生产 PG 落库与 BYOK 归属随 M4b;
|
||||||
|
- 健康面 `GET /healthz` `GET /readyz`;服务令牌 `SHENSHI_SERVICE_TOKEN`(开发档可空)。
|
||||||
|
|
||||||
## 目录结构
|
## 目录结构
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,88 @@
|
|||||||
|
// Package metering 用量计量:每次尝试一条记录(stdout 结构化日志 + 进程内聚合)。
|
||||||
|
// 生产部署接 PG(shenshi-gateway.md §7);M4a 先以日志与内存聚合支撑 /me/usage 与对账抽样。
|
||||||
|
package metering
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"log"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Attempt 一次端点尝试的计量记录。
|
||||||
|
type Attempt struct {
|
||||||
|
RunID string `json:"run_id"`
|
||||||
|
CharacterID string `json:"character_id,omitempty"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
Slot string `json:"slot"`
|
||||||
|
Endpoint string `json:"endpoint"`
|
||||||
|
Status string `json:"status"` // ok / error / first_byte_timeout
|
||||||
|
ErrorCode string `json:"error_code,omitempty"`
|
||||||
|
FirstTokenMS int64 `json:"first_token_ms"`
|
||||||
|
TotalMS int64 `json:"total_ms"`
|
||||||
|
PromptTokens int `json:"prompt_tokens"`
|
||||||
|
CompletionTokens int `json:"completion_tokens"`
|
||||||
|
At time.Time `json:"at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// UsageAgg /me/usage 聚合行。
|
||||||
|
type UsageAgg struct {
|
||||||
|
Model string `json:"model"`
|
||||||
|
Endpoint string `json:"endpoint"`
|
||||||
|
Requests int `json:"requests"`
|
||||||
|
OK int `json:"ok"`
|
||||||
|
PromptTokens int `json:"prompt_tokens"`
|
||||||
|
CompletionTokens int `json:"completion_tokens"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Log 进程内计量器(有界环形,避免无界增长)。
|
||||||
|
type Log struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
entries []Attempt
|
||||||
|
limit int
|
||||||
|
}
|
||||||
|
|
||||||
|
// New 构造计量器(保留最近 N 条)。
|
||||||
|
func New(limit int) *Log {
|
||||||
|
return &Log{entries: make([]Attempt, 0, limit), limit: limit}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Record 记录一次尝试:stdout JSON 日志 + 内存追加。
|
||||||
|
func (l *Log) Record(a Attempt) {
|
||||||
|
a.At = time.Now().UTC()
|
||||||
|
if raw, err := json.Marshal(a); err == nil {
|
||||||
|
log.Printf(`{"kind":"shenshi_usage",%s`, string(raw[1:]))
|
||||||
|
}
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
if len(l.entries) >= l.limit {
|
||||||
|
l.entries = l.entries[1:]
|
||||||
|
}
|
||||||
|
l.entries = append(l.entries, a)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Aggregate 按 模型×端点 聚合。
|
||||||
|
func (l *Log) Aggregate() []UsageAgg {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
agg := map[string]*UsageAgg{}
|
||||||
|
for _, e := range l.entries {
|
||||||
|
key := e.Model + "|" + e.Endpoint
|
||||||
|
row := agg[key]
|
||||||
|
if row == nil {
|
||||||
|
row = &UsageAgg{Model: e.Model, Endpoint: e.Endpoint}
|
||||||
|
agg[key] = row
|
||||||
|
}
|
||||||
|
row.Requests++
|
||||||
|
if e.Status == "ok" {
|
||||||
|
row.OK++
|
||||||
|
row.PromptTokens += e.PromptTokens
|
||||||
|
row.CompletionTokens += e.CompletionTokens
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out := make([]UsageAgg, 0, len(agg))
|
||||||
|
for _, row := range agg {
|
||||||
|
out = append(out, *row)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
@@ -0,0 +1,75 @@
|
|||||||
|
// 候选路由与冷却:凭证级(401/403)与绑定级(429/5xx/超时)冷却,粘性会话与故障迭代。
|
||||||
|
package route
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// 冷却时长:凭证错误长冷却,限流/超时短冷却(避免整渠道误杀,shenshi-gateway.md §4)。
|
||||||
|
const (
|
||||||
|
credentialCooldown = 10 * time.Minute
|
||||||
|
bindingCooldown = 30 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// Cooldowns 进程内冷却状态:key 形如 "cred:<endpoint>:<keyhash>" / "bind:<endpoint>:<model>"。
|
||||||
|
type Cooldowns struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
until map[string]time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewCooldowns 构造冷却表。
|
||||||
|
func NewCooldowns() *Cooldowns {
|
||||||
|
return &Cooldowns{until: make(map[string]time.Time)}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Mark 凭证级(credential=true)或绑定级冷却。
|
||||||
|
func (c *Cooldowns) Mark(key string, credential bool) {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
duration := bindingCooldown
|
||||||
|
if credential {
|
||||||
|
duration = credentialCooldown
|
||||||
|
}
|
||||||
|
c.until[key] = time.Now().Add(duration)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Active 该冷却键当前是否生效。
|
||||||
|
func (c *Cooldowns) Active(key string) bool {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
t, ok := c.until[key]
|
||||||
|
return ok && time.Now().Before(t)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sticky 粘性会话:同一 run 连续请求优先同端点(进程内,run 生命周期短)。
|
||||||
|
type Sticky struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
binds map[string]string
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSticky 构造粘性表。
|
||||||
|
func NewSticky() *Sticky {
|
||||||
|
return &Sticky{binds: make(map[string]string)}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remember 记录 run → 端点绑定。
|
||||||
|
func (s *Sticky) Remember(runID, endpointName string) {
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
s.binds[runID] = endpointName
|
||||||
|
}
|
||||||
|
|
||||||
|
// Recall 取 run 的粘性端点(无则空串)。
|
||||||
|
func (s *Sticky) Recall(runID string) string {
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
return s.binds[runID]
|
||||||
|
}
|
||||||
|
|
||||||
|
// Forget run 结束后清理。
|
||||||
|
func (s *Sticky) Forget(runID string) {
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
delete(s.binds, runID)
|
||||||
|
}
|
||||||
@@ -1,46 +1,150 @@
|
|||||||
// Package route 维护逻辑模型/能力槽位到候选端点的路由注册表。
|
// Package route 维护逻辑模型/能力槽位到候选端点的路由注册表。
|
||||||
// M0 为单一 stub 端点的静态表;候选迭代、健康冷却与故障切换按
|
// 端点来源:环境 SHENSHI_ENDPOINTS(JSON 数组);未配置时回落 stub(骨架形态)。
|
||||||
// docs/architecture/shenshi-gateway.md §4 于 M4 落地。
|
// 冷却粒度(shenshi-gateway.md §4):401/403 只冷却凭证;429/5xx/超时只冷却该端点绑定。
|
||||||
package route
|
package route
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"sort"
|
||||||
|
"time"
|
||||||
|
|
||||||
"hunjing/shenshi/internal/transform"
|
"hunjing/shenshi/internal/transform"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Endpoint 一个可路由的出站端点:类型 + 适配器实例。
|
// Endpoint 一个可路由的出站端点。
|
||||||
type Endpoint struct {
|
type Endpoint struct {
|
||||||
Type string
|
Name string
|
||||||
|
Type string // openai_compatible / local / stub(契约词典 §14.2)
|
||||||
Adaptor transform.Adaptor
|
Adaptor transform.Adaptor
|
||||||
Priority int
|
Priority int
|
||||||
|
// 凭证指纹(凭证级冷却键);stub/openai 兼容端点各自填充
|
||||||
|
CredentialKey string
|
||||||
|
// 绑定级冷却键的模型名部分
|
||||||
|
UpstreamModel string
|
||||||
|
}
|
||||||
|
|
||||||
|
// EndpointSpec 环境配置形态。
|
||||||
|
type EndpointSpec struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
BaseURL string `json:"base_url"`
|
||||||
|
APIKey string `json:"api_key"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
Priority int `json:"priority"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// Registry 槽位 → 候选端点列表(按 Priority 升序尝试)。
|
// Registry 槽位 → 候选端点列表(按 Priority 升序尝试)。
|
||||||
type Registry struct {
|
type Registry struct {
|
||||||
bySlot map[string][]Endpoint
|
bySlot map[string][]Endpoint
|
||||||
|
// Cooldowns/FirstByteTimeoutMS 由 server 层共享使用
|
||||||
|
Cooldowns *Cooldowns
|
||||||
|
Sticky *Sticky
|
||||||
|
FirstByteTimeout time.Duration
|
||||||
}
|
}
|
||||||
|
|
||||||
// New 构建注册表并登记默认端点。
|
// New 构建注册表:读 SHENSHI_ENDPOINTS(JSON),无配置时登记 stub。
|
||||||
func New() *Registry {
|
func New() *Registry {
|
||||||
r := &Registry{bySlot: make(map[string][]Endpoint)}
|
r := &Registry{
|
||||||
r.Register("chat", Endpoint{Type: "stub", Adaptor: transform.StubAdaptor{}, Priority: 100})
|
bySlot: make(map[string][]Endpoint),
|
||||||
|
Cooldowns: NewCooldowns(),
|
||||||
|
Sticky: NewSticky(),
|
||||||
|
FirstByteTimeout: 10 * time.Second,
|
||||||
|
}
|
||||||
|
specs := loadSpecs()
|
||||||
|
if len(specs) == 0 {
|
||||||
|
r.Register("chat", Endpoint{Type: "stub", Adaptor: transform.StubAdaptor{}, Priority: 100, UpstreamModel: "stub"})
|
||||||
|
return r
|
||||||
|
}
|
||||||
|
for _, spec := range specs {
|
||||||
|
if spec.Type != "openai_compatible" && spec.Type != "local" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
adaptor := &transform.OpenAICompatAdaptor{
|
||||||
|
EndpointName: spec.Name,
|
||||||
|
BaseURL: spec.BaseURL,
|
||||||
|
APIKey: spec.APIKey,
|
||||||
|
UpstreamModel: spec.Model,
|
||||||
|
}
|
||||||
|
cred := spec.Name
|
||||||
|
if spec.APIKey != "" {
|
||||||
|
cred = spec.Name + ":" + fmt.Sprintf("%x", len(spec.APIKey)) // 不落密钥原文,仅指纹
|
||||||
|
}
|
||||||
|
r.Register("chat", Endpoint{
|
||||||
|
Name: spec.Name, Type: spec.Type, Adaptor: adaptor,
|
||||||
|
Priority: spec.Priority, CredentialKey: cred, UpstreamModel: spec.Model,
|
||||||
|
})
|
||||||
|
}
|
||||||
return r
|
return r
|
||||||
}
|
}
|
||||||
|
|
||||||
// Register 登记端点到指定能力槽位(契约词典 §14.1 model_slot 取值)。
|
func loadSpecs() []EndpointSpec {
|
||||||
|
raw := os.Getenv("SHENSHI_ENDPOINTS")
|
||||||
|
if raw == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var specs []EndpointSpec
|
||||||
|
if err := json.Unmarshal([]byte(raw), &specs); err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return specs
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reset 清空槽位(测试装配用)。
|
||||||
|
func (r *Registry) Reset(slot string) { delete(r.bySlot, slot) }
|
||||||
|
|
||||||
|
// Register 登记端点到指定能力槽位(契约词典 §14.1)。
|
||||||
func (r *Registry) Register(slot string, ep Endpoint) {
|
func (r *Registry) Register(slot string, ep Endpoint) {
|
||||||
r.bySlot[slot] = append(r.bySlot[slot], ep)
|
r.bySlot[slot] = append(r.bySlot[slot], ep)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Resolve 返回槽位的候选端点;无候选即错误(HJ-2001 语义)。
|
// Candidates 返回槽位候选(Priority 升序;粘性端点置首;冷却中端点排后不剔除——
|
||||||
func (r *Registry) Resolve(slot string) ([]Endpoint, error) {
|
// 全部冷却时仍可尝试最优先者,避免死锁)。
|
||||||
|
func (r *Registry) Candidates(slot, runID string) ([]Endpoint, error) {
|
||||||
if slot == "" {
|
if slot == "" {
|
||||||
slot = "chat"
|
slot = "chat"
|
||||||
}
|
}
|
||||||
eps := r.bySlot[slot]
|
eps := append([]Endpoint(nil), r.bySlot[slot]...)
|
||||||
if len(eps) == 0 {
|
if len(eps) == 0 {
|
||||||
return nil, fmt.Errorf("槽位 %s 无可用端点", slot)
|
return nil, fmt.Errorf("槽位 %s 无可用端点", slot)
|
||||||
}
|
}
|
||||||
|
sort.SliceStable(eps, func(i, j int) bool { return eps[i].Priority < eps[j].Priority })
|
||||||
|
if runID != "" {
|
||||||
|
if stickyName := r.Sticky.Recall(runID); stickyName != "" {
|
||||||
|
for i, ep := range eps {
|
||||||
|
if ep.Name == stickyName {
|
||||||
|
eps = append([]Endpoint{ep}, append(eps[:i:i], eps[i+1:]...)...)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// 冷却中的端点移到队尾(保留兜底能力)
|
||||||
|
cooling := 0
|
||||||
|
for i := 0; i < len(eps)-cooling; {
|
||||||
|
if r.Cooldowns.Active(bindingKey(eps[i])) || r.Cooldowns.Active(credKey(eps[i])) {
|
||||||
|
ep := eps[i]
|
||||||
|
eps = append(append(eps[:i:i], eps[i+1:]...), ep)
|
||||||
|
cooling++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
i++
|
||||||
|
}
|
||||||
return eps, nil
|
return eps, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func bindingKey(ep Endpoint) string {
|
||||||
|
return "bind:" + ep.Name + ":" + ep.UpstreamModel
|
||||||
|
}
|
||||||
|
|
||||||
|
func credKey(ep Endpoint) string {
|
||||||
|
if ep.CredentialKey == "" {
|
||||||
|
return "cred:" + ep.Name
|
||||||
|
}
|
||||||
|
return "cred:" + ep.CredentialKey
|
||||||
|
}
|
||||||
|
|
||||||
|
// BindingCooldownKey / CredentialCooldownKey 供 server 层标记冷却。
|
||||||
|
func BindingCooldownKey(ep Endpoint) string { return bindingKey(ep) }
|
||||||
|
func CredentialCooldownKey(ep Endpoint) string { return credKey(ep) }
|
||||||
@@ -0,0 +1,161 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"hunjing/shenshi/internal/route"
|
||||||
|
"hunjing/shenshi/internal/transform"
|
||||||
|
)
|
||||||
|
|
||||||
|
// fakeOpenAI 构造一个 OpenAI 兼容 SSE 假上游:可控延迟/状态码/内容。
|
||||||
|
func fakeOpenAI(t *testing.T, delay time.Duration, status int, content string) *httptest.Server {
|
||||||
|
t.Helper()
|
||||||
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
if status != 0 {
|
||||||
|
w.WriteHeader(status)
|
||||||
|
_, _ = w.Write([]byte(`{"error":{"message":"boom"}}`))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if delay > 0 {
|
||||||
|
time.Sleep(delay)
|
||||||
|
}
|
||||||
|
w.Header().Set("content-type", "text/event-stream")
|
||||||
|
flusher := w.(http.Flusher)
|
||||||
|
_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"" + content + "\"}}]}\n\n"))
|
||||||
|
flusher.Flush()
|
||||||
|
_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":4}}\n\n"))
|
||||||
|
flusher.Flush()
|
||||||
|
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
func registryWith(t *testing.T, firstByte time.Duration, eps ...route.Endpoint) *route.Registry {
|
||||||
|
t.Helper()
|
||||||
|
r := route.New()
|
||||||
|
r.FirstByteTimeout = firstByte
|
||||||
|
// 清掉默认 stub,仅保留测试端点
|
||||||
|
r.Reset("chat")
|
||||||
|
for _, ep := range eps {
|
||||||
|
r.Register("chat", ep)
|
||||||
|
}
|
||||||
|
return r
|
||||||
|
}
|
||||||
|
|
||||||
|
func oaEndpoint(name, baseURL string, priority int) route.Endpoint {
|
||||||
|
return route.Endpoint{
|
||||||
|
Name: name, Type: "openai_compatible", Priority: priority,
|
||||||
|
Adaptor: &transform.OpenAICompatAdaptor{EndpointName: name, BaseURL: baseURL, UpstreamModel: "test-model"},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func foPostChat(ts *httptest.Server, body string) (*http.Response, string) {
|
||||||
|
resp, err := ts.Client().Post(ts.URL+"/v1/chat/completions", "application/json", strings.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
return nil, ""
|
||||||
|
}
|
||||||
|
var sb strings.Builder
|
||||||
|
scanner := bufio.NewScanner(resp.Body)
|
||||||
|
for scanner.Scan() {
|
||||||
|
sb.WriteString(scanner.Text())
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
return resp, sb.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func foChatBody(run string) string {
|
||||||
|
return `{"ir_version":1,"slot":"chat","messages":[{"role":"user","content":"你好"}],"run_id":"` + run + `","stream":true}`
|
||||||
|
}
|
||||||
|
|
||||||
|
// 场景:首端点首字超时 → 切道到次端点成功(不向客户端暴露失败)。
|
||||||
|
func TestFailoverOnFirstByteTimeout(t *testing.T) {
|
||||||
|
slow := fakeOpenAI(t, 500*time.Millisecond, 0, "慢端点")
|
||||||
|
fast := fakeOpenAI(t, 0, 0, "快端点")
|
||||||
|
defer slow.Close()
|
||||||
|
defer fast.Close()
|
||||||
|
|
||||||
|
srv := New(Config{}, registryWith(t, 50*time.Millisecond, oaEndpoint("slow", slow.URL, 10), oaEndpoint("fast", fast.URL, 20)))
|
||||||
|
ts := httptest.NewServer(srv)
|
||||||
|
defer ts.Close()
|
||||||
|
|
||||||
|
resp, body := foPostChat(ts, foChatBody("run_fo"))
|
||||||
|
if resp == nil || resp.StatusCode != http.StatusOK {
|
||||||
|
t.Fatalf("状态码异常:%v", resp)
|
||||||
|
}
|
||||||
|
if !strings.Contains(body, "快端点") || strings.Contains(body, "慢端点") {
|
||||||
|
t.Fatalf("应切道到快端点:%s", body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(body, `"finish_reason":"stop"`) || !strings.Contains(body, "data: [DONE]") {
|
||||||
|
t.Fatalf("缺少终帧/DONE:%s", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 场景:首端点 401 → 凭证级冷却;下一请求直接跳过它走次端点。
|
||||||
|
func TestCredentialCooldownSkipsEndpoint(t *testing.T) {
|
||||||
|
bad := fakeOpenAI(t, 0, http.StatusUnauthorized, "")
|
||||||
|
good := fakeOpenAI(t, 0, 0, "好端点")
|
||||||
|
defer bad.Close()
|
||||||
|
defer good.Close()
|
||||||
|
|
||||||
|
srv := New(Config{}, registryWith(t, time.Second, oaEndpoint("bad", bad.URL, 10), oaEndpoint("good", good.URL, 20)))
|
||||||
|
ts := httptest.NewServer(srv)
|
||||||
|
defer ts.Close()
|
||||||
|
|
||||||
|
// 第一次:bad 401 → 冷却 → 切 good
|
||||||
|
resp, body := foPostChat(ts, foChatBody("run_a"))
|
||||||
|
if resp == nil || resp.StatusCode != http.StatusOK || !strings.Contains(body, "好端点") {
|
||||||
|
t.Fatalf("首次应切到好端点:%v %s", resp, body)
|
||||||
|
}
|
||||||
|
// 第二次:bad 仍在冷却,直接命中 good(若未跳过会再花一次 401 往返——这里断言成功即可)
|
||||||
|
resp2, body2 := foPostChat(ts, foChatBody("run_b"))
|
||||||
|
if resp2 == nil || resp2.StatusCode != http.StatusOK || !strings.Contains(body2, "好端点") {
|
||||||
|
t.Fatalf("冷却后应直达好端点:%v", resp2)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 场景:全部候选失败且未写字节 → HJ-2001 JSON(不是半截 SSE)。
|
||||||
|
func TestAllCandidatesFailedReturnsJsonError(t *testing.T) {
|
||||||
|
bad := fakeOpenAI(t, 0, http.StatusBadGateway, "")
|
||||||
|
defer bad.Close()
|
||||||
|
srv := New(Config{}, registryWith(t, time.Second, oaEndpoint("bad", bad.URL, 10)))
|
||||||
|
ts := httptest.NewServer(srv)
|
||||||
|
defer ts.Close()
|
||||||
|
|
||||||
|
resp, body := foPostChat(ts, foChatBody("run_x"))
|
||||||
|
if resp.StatusCode != http.StatusServiceUnavailable {
|
||||||
|
t.Fatalf("状态码 = %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
if !strings.Contains(body, "HJ-2001") || strings.Contains(body, "text/event-stream") {
|
||||||
|
t.Fatalf("应为 JSON 错误而非 SSE:%s", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 场景:用量聚合端点返回 ok 计数与 token。
|
||||||
|
func TestUsageAggregate(t *testing.T) {
|
||||||
|
good := fakeOpenAI(t, 0, 0, "内容")
|
||||||
|
defer good.Close()
|
||||||
|
srv := New(Config{}, registryWith(t, time.Second, oaEndpoint("good", good.URL, 10)))
|
||||||
|
ts := httptest.NewServer(srv)
|
||||||
|
defer ts.Close()
|
||||||
|
|
||||||
|
if resp, _ := foPostChat(ts, foChatBody("run_u")); resp == nil || resp.StatusCode != 200 {
|
||||||
|
t.Fatal("请求失败")
|
||||||
|
}
|
||||||
|
usageResp, err := ts.Client().Get(ts.URL + "/me/usage")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var agg []map[string]any
|
||||||
|
defer func() { _ = usageResp.Body.Close() }()
|
||||||
|
if err := json.NewDecoder(usageResp.Body).Decode(&agg); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(agg) != 1 || agg[0]["ok"].(float64) != 1 || agg[0]["prompt_tokens"].(float64) != 3 {
|
||||||
|
t.Fatalf("聚合不符:%v", agg)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,33 +1,42 @@
|
|||||||
// Package server 神识 HTTP 服务:数据面(IR 入站 → SSE 出站)与健康面。
|
// Package server 神识 HTTP 服务:数据面(IR 入站 → SSE 出站)、健康面与用量面。
|
||||||
|
//
|
||||||
|
// 故障切换红线:首帧到达前不提交 SSE 响应——失败可整体换道并返回 JSON 错误;
|
||||||
|
// 已写出字节后不再换道,如实以 error 终帧收尾(shenshi-gateway.md §4/§5)。
|
||||||
package server
|
package server
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
"time"
|
||||||
|
|
||||||
"hunjing/shenshi/internal/ir"
|
"hunjing/shenshi/internal/ir"
|
||||||
|
"hunjing/shenshi/internal/metering"
|
||||||
"hunjing/shenshi/internal/relay"
|
"hunjing/shenshi/internal/relay"
|
||||||
"hunjing/shenshi/internal/route"
|
"hunjing/shenshi/internal/route"
|
||||||
|
"hunjing/shenshi/internal/transform"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Config 服务配置。
|
// Config 服务配置。
|
||||||
type Config struct {
|
type Config struct {
|
||||||
ServiceToken string // 灵台→神识服务令牌;为空时开发档不校验
|
ServiceToken string
|
||||||
}
|
}
|
||||||
|
|
||||||
// Server 神识 HTTP 服务。
|
// Server 神识 HTTP 服务。
|
||||||
type Server struct {
|
type Server struct {
|
||||||
cfg Config
|
cfg Config
|
||||||
routes *route.Registry
|
routes *route.Registry
|
||||||
|
usage *metering.Log
|
||||||
mux *http.ServeMux
|
mux *http.ServeMux
|
||||||
}
|
}
|
||||||
|
|
||||||
// New 构建服务并注册路由。
|
// New 构建服务并注册路由。
|
||||||
func New(cfg Config, routes *route.Registry) *Server {
|
func New(cfg Config, routes *route.Registry) *Server {
|
||||||
s := &Server{cfg: cfg, routes: routes, mux: http.NewServeMux()}
|
s := &Server{cfg: cfg, routes: routes, usage: metering.New(10_000), mux: http.NewServeMux()}
|
||||||
s.mux.HandleFunc("GET /healthz", s.handleHealth)
|
s.mux.HandleFunc("GET /healthz", s.handleHealth)
|
||||||
s.mux.HandleFunc("GET /readyz", s.handleHealth)
|
s.mux.HandleFunc("GET /readyz", s.handleHealth)
|
||||||
s.mux.Handle("POST /v1/chat/completions", s.auth(http.HandlerFunc(s.handleChat)))
|
s.mux.Handle("POST /v1/chat/completions", s.auth(http.HandlerFunc(s.handleChat)))
|
||||||
|
s.mux.Handle("GET /me/usage", s.auth(http.HandlerFunc(s.handleUsage)))
|
||||||
return s
|
return s
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -41,8 +50,7 @@ func (s *Server) auth(next http.Handler) http.Handler {
|
|||||||
next.ServeHTTP(w, r)
|
next.ServeHTTP(w, r)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
credential := r.Header.Get("Authorization")
|
if r.Header.Get("Authorization") != "Bearer "+s.cfg.ServiceToken {
|
||||||
if credential != "Bearer "+s.cfg.ServiceToken {
|
|
||||||
writeError(w, http.StatusUnauthorized, "HJ-2003", "服务令牌无效")
|
writeError(w, http.StatusUnauthorized, "HJ-2003", "服务令牌无效")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@@ -50,7 +58,7 @@ func (s *Server) auth(next http.Handler) http.Handler {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
// handleChat 数据面:解析 IR → 路由 → 适配器流出 SSE。
|
// handleChat 数据面:IR → 候选迭代(首字超时/冷却标记/粘性)→ SSE。
|
||||||
func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
|
func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||||
var req ir.ChatRequest
|
var req ir.ChatRequest
|
||||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||||
@@ -61,35 +69,136 @@ func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
|
|||||||
writeError(w, http.StatusBadRequest, "HJ-2002", err.Error())
|
writeError(w, http.StatusBadRequest, "HJ-2002", err.Error())
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
endpoints, err := s.routes.Resolve(req.Slot)
|
candidates, err := s.routes.Candidates(req.Slot, req.RunID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
writeError(w, http.StatusServiceUnavailable, "HJ-2001", err.Error())
|
writeError(w, http.StatusServiceUnavailable, "HJ-2001", err.Error())
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
relay.SSEHeaders(w)
|
var lastErr error
|
||||||
w.WriteHeader(http.StatusOK)
|
for _, endpoint := range candidates {
|
||||||
|
res := s.tryEndpoint(w, r, endpoint, &req)
|
||||||
|
s.usage.Record(res.Attempt)
|
||||||
|
if res.Committed {
|
||||||
|
// 已向客户端写出:无论成败都不能再换道
|
||||||
|
if res.Status == "ok" {
|
||||||
|
s.routes.Sticky.Remember(req.RunID, endpoint.Name)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// 未提交:分类冷却后换下一候选
|
||||||
|
if res.Upstream != nil {
|
||||||
|
if res.Upstream.StatusCode == 401 || res.Upstream.StatusCode == 403 {
|
||||||
|
s.routes.Cooldowns.Mark(route.CredentialCooldownKey(endpoint), true)
|
||||||
|
} else {
|
||||||
|
s.routes.Cooldowns.Mark(route.BindingCooldownKey(endpoint), false)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
s.routes.Cooldowns.Mark(route.BindingCooldownKey(endpoint), false)
|
||||||
|
}
|
||||||
|
lastErr = res.Err
|
||||||
|
}
|
||||||
|
writeError(w, http.StatusServiceUnavailable, "HJ-2001", "全部候选端点失败:"+errString(lastErr))
|
||||||
|
}
|
||||||
|
|
||||||
|
type attempt struct {
|
||||||
|
metering.Attempt
|
||||||
|
Committed bool // 是否已向客户端写出 SSE
|
||||||
|
Upstream *transform.UpstreamError // 上游分类(可空)
|
||||||
|
Err error // 原始错误
|
||||||
|
}
|
||||||
|
|
||||||
|
func errString(err error) string {
|
||||||
|
if err == nil {
|
||||||
|
return "未知错误"
|
||||||
|
}
|
||||||
|
return err.Error()
|
||||||
|
}
|
||||||
|
|
||||||
|
// tryEndpoint 单端点尝试。
|
||||||
|
func (s *Server) tryEndpoint(w http.ResponseWriter, r *http.Request, endpoint route.Endpoint, req *ir.ChatRequest) attempt {
|
||||||
|
start := time.Now()
|
||||||
|
ctx, cancel := context.WithTimeout(r.Context(), s.routes.FirstByteTimeout)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
deltas := make(chan ir.Delta, 16)
|
deltas := make(chan ir.Delta, 16)
|
||||||
streamErr := make(chan error, 1)
|
errCh := make(chan error, 1)
|
||||||
go func() {
|
go func() {
|
||||||
defer close(deltas)
|
defer close(deltas)
|
||||||
streamErr <- endpoints[0].Adaptor.Stream(r.Context(), &req, deltas)
|
errCh <- endpoint.Adaptor.Stream(ctx, req, deltas)
|
||||||
}()
|
}()
|
||||||
|
|
||||||
|
res := attempt{Attempt: metering.Attempt{
|
||||||
|
RunID: req.RunID, CharacterID: req.CharacterID,
|
||||||
|
Model: req.Model, Slot: req.Slot, Endpoint: endpoint.Name,
|
||||||
|
}}
|
||||||
|
|
||||||
|
// 阶段一:等首帧或失败——此阶段未提交响应,可整体回退
|
||||||
|
var first ir.Delta
|
||||||
|
select {
|
||||||
|
case d, ok := <-deltas:
|
||||||
|
if !ok {
|
||||||
|
res.Err = <-errCh
|
||||||
|
res.Upstream, _ = res.Err.(*transform.UpstreamError)
|
||||||
|
res.Status, res.TotalMS = failStatus(res.Err), time.Since(start).Milliseconds()
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
first = d
|
||||||
|
case <-ctx.Done():
|
||||||
|
res.Status, res.Err = "first_byte_timeout", ctx.Err()
|
||||||
|
res.TotalMS = time.Since(start).Milliseconds()
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
|
// 阶段二:首帧已到——提交 SSE,流式转发
|
||||||
|
relay.SSEHeaders(w)
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
res.Committed = true
|
||||||
|
res.FirstTokenMS = time.Since(start).Milliseconds()
|
||||||
|
_ = relay.WriteDelta(w, first)
|
||||||
|
|
||||||
for d := range deltas {
|
for d := range deltas {
|
||||||
|
if d.Delta != "" || d.FinishReason != "" {
|
||||||
if err := relay.WriteDelta(w, d); err != nil {
|
if err := relay.WriteDelta(w, d); err != nil {
|
||||||
return // 客户端断开:取消上游由 r.Context() 传播
|
break // 客户端断开
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if err := <-streamErr; err != nil {
|
if d.FinishReason != "" {
|
||||||
// 流已开始,只能如实中断:写一条错误帧后终止(不伪装 [DONE])
|
res.Attempt.PromptTokens = d.UsagePromptTokens
|
||||||
_ = relay.WriteDelta(w, ir.Delta{FinishReason: "error", Delta: ""})
|
res.Attempt.CompletionTokens = d.UsageCompletionTokens
|
||||||
|
}
|
||||||
|
}
|
||||||
|
err := <-errCh
|
||||||
|
res.TotalMS = time.Since(start).Milliseconds()
|
||||||
|
if err == nil {
|
||||||
_ = relay.WriteDone(w)
|
_ = relay.WriteDone(w)
|
||||||
return
|
res.Status = "ok"
|
||||||
|
return res
|
||||||
}
|
}
|
||||||
if err := relay.WriteDone(w); err != nil {
|
// 流中失败:如实以 error 终帧收尾,不伪装完成
|
||||||
return
|
res.Err = err
|
||||||
|
res.Upstream, _ = err.(*transform.UpstreamError)
|
||||||
|
res.Status = "error"
|
||||||
|
res.ErrorCode = "stream_written"
|
||||||
|
_ = relay.WriteDelta(w, ir.Delta{FinishReason: "error"})
|
||||||
|
_ = relay.WriteDone(w)
|
||||||
|
return res
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func failStatus(err error) string {
|
||||||
|
if up, ok := err.(*transform.UpstreamError); ok {
|
||||||
|
if up.StatusCode == 0 {
|
||||||
|
return "timeout_or_network"
|
||||||
|
}
|
||||||
|
return "upstream_error"
|
||||||
|
}
|
||||||
|
return "error"
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleUsage 用量聚合(/me/usage;BYOK 归属维度随 M4b 接入)。
|
||||||
|
func (s *Server) handleUsage(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_ = json.NewEncoder(w).Encode(s.usage.Aggregate())
|
||||||
}
|
}
|
||||||
|
|
||||||
// handleHealth 健康面。
|
// handleHealth 健康面。
|
||||||
|
|||||||
@@ -0,0 +1,138 @@
|
|||||||
|
// 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()
|
||||||
|
}
|
||||||
Reference in new issue
Block a user