重构项目结构并纳入版本管理
- 源码移至 src/go/，编译产物输出到 build/（不入库） - 新增 .gitignore：排除 data/、build/、*.exe、*.log、说明.md - 新增 build.ps1：从 src/go 编译到 build/relay.exe - 移除 说明.md 的版本跟踪（含上游密钥，仅保留在本地） Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
185
src/go/sse.go
Normal file
185
src/go/sse.go
Normal file
@@ -0,0 +1,185 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// streamRecorder 包在 resp.Body 外层：数据被客户端读取时同步流过，
|
||||
// 一边原样转发（客户端实时收到每个 chunk），一边把全量原文拼回 + 解析 SSE usage。
|
||||
// 读到 EOF（流结束）时回调 onDone，由它落库。
|
||||
//
|
||||
// 这是流式代理的标准做法：用一个 io.ReadCloser 包装上游 body，
|
||||
// 不做 io.ReadAll，避免缓冲整条流导致客户端要等所有 chunk 到齐才收到第一个字。
|
||||
type streamRecorder struct {
|
||||
src io.ReadCloser
|
||||
raw bytes.Buffer // 全量原文（拼回所有 chunk）
|
||||
parser *sseParser
|
||||
onDone func(raw []byte, in, out int, model string)
|
||||
done bool
|
||||
}
|
||||
|
||||
func newStreamRecorder(src io.ReadCloser, format string, onDone func(raw []byte, in, out int, model string)) *streamRecorder {
|
||||
return &streamRecorder{
|
||||
src: src,
|
||||
parser: newSSEParser(format),
|
||||
onDone: onDone,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *streamRecorder) Read(p []byte) (int, error) {
|
||||
n, err := s.src.Read(p)
|
||||
if n > 0 {
|
||||
// 原样留一份，并喂给 SSE 解析器累积 usage
|
||||
s.raw.Write(p[:n])
|
||||
s.parser.feed(p[:n])
|
||||
}
|
||||
if err == io.EOF && !s.done {
|
||||
s.done = true
|
||||
s.parser.flush()
|
||||
in, out, model := s.parser.result()
|
||||
s.onDone(s.raw.Bytes(), in, out, model)
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (s *streamRecorder) Close() error {
|
||||
// 客户端提前断开等情况：若还没触发 onDone，这里兜底落库一次。
|
||||
if !s.done {
|
||||
s.done = true
|
||||
s.parser.flush()
|
||||
in, out, model := s.parser.result()
|
||||
s.onDone(s.raw.Bytes(), in, out, model)
|
||||
}
|
||||
return s.src.Close()
|
||||
}
|
||||
|
||||
// sseParser 按行累积 SSE 数据，跨 chunk 解析 OpenAI / Claude 两家流式 usage。
|
||||
//
|
||||
// OpenAI 流式：usage 出现在末尾带 stream_options.include_usage 的 chunk 里，
|
||||
// 形如 data: {...,"usage":{"prompt_tokens":..,"completion_tokens":..}}；
|
||||
// 没开 include_usage 则 usage 落空（如实记 0）。
|
||||
// Claude 流式：message_start 事件给 input_tokens（usage.input_tokens），
|
||||
// message_delta 事件累加 output_tokens（usage.output_tokens），需跨事件攒。
|
||||
type sseParser struct {
|
||||
format string
|
||||
buf bytes.Buffer // 未成行的残留字节
|
||||
|
||||
in int
|
||||
out int
|
||||
model string
|
||||
}
|
||||
|
||||
func newSSEParser(format string) *sseParser {
|
||||
return &sseParser{format: format}
|
||||
}
|
||||
|
||||
// feed 接收任意切片的字节，按行切出完整的 SSE 行交给 parseLine。
|
||||
func (p *sseParser) feed(b []byte) {
|
||||
p.buf.Write(b)
|
||||
data := p.buf.Bytes()
|
||||
for {
|
||||
i := bytes.IndexByte(data, '\n')
|
||||
if i < 0 {
|
||||
break
|
||||
}
|
||||
line := data[:i]
|
||||
p.parseLine(line)
|
||||
data = data[i+1:]
|
||||
}
|
||||
// 保留未成行的尾部
|
||||
rest := make([]byte, len(data))
|
||||
copy(rest, data)
|
||||
p.buf.Reset()
|
||||
p.buf.Write(rest)
|
||||
}
|
||||
|
||||
// flush 处理流结束时缓冲里残留的最后一行（可能没有结尾换行）。
|
||||
func (p *sseParser) flush() {
|
||||
if p.buf.Len() > 0 {
|
||||
p.parseLine(p.buf.Bytes())
|
||||
p.buf.Reset()
|
||||
}
|
||||
}
|
||||
|
||||
// parseLine 解析单行 SSE。只关心 "data:" 行里的 JSON。
|
||||
func (p *sseParser) parseLine(line []byte) {
|
||||
s := strings.TrimRight(string(line), "\r")
|
||||
s = strings.TrimSpace(s)
|
||||
if !strings.HasPrefix(s, "data:") {
|
||||
return
|
||||
}
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(s, "data:"))
|
||||
if payload == "" || payload == "[DONE]" {
|
||||
return
|
||||
}
|
||||
|
||||
var m map[string]json.RawMessage
|
||||
if err := json.Unmarshal([]byte(payload), &m); err != nil {
|
||||
return // 非 JSON 行，跳过
|
||||
}
|
||||
|
||||
// model 字段（两家都有，取到就记）
|
||||
if raw, ok := m["model"]; ok {
|
||||
var mod string
|
||||
if json.Unmarshal(raw, &mod) == nil && mod != "" {
|
||||
p.model = mod
|
||||
}
|
||||
}
|
||||
|
||||
// usage 可能直接在顶层（OpenAI 末尾 chunk），
|
||||
// 也可能在 message.usage 里（Claude message_start）。
|
||||
if raw, ok := m["usage"]; ok {
|
||||
p.applyUsage(raw)
|
||||
}
|
||||
if raw, ok := m["message"]; ok {
|
||||
var msg map[string]json.RawMessage
|
||||
if json.Unmarshal(raw, &msg) == nil {
|
||||
if u, ok := msg["usage"]; ok {
|
||||
p.applyUsage(u)
|
||||
}
|
||||
if mraw, ok := msg["model"]; ok {
|
||||
var mod string
|
||||
if json.Unmarshal(mraw, &mod) == nil && mod != "" {
|
||||
p.model = mod
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// applyUsage 从一段 usage JSON 里抽 token 数，兼容两家命名。
|
||||
// Claude 的 output_tokens 在 message_delta 里是累计值（取最大即可），
|
||||
// input_tokens 在 message_start 给定后不再变；OpenAI 末尾一次性给全量。
|
||||
func (p *sseParser) applyUsage(raw json.RawMessage) {
|
||||
var u struct {
|
||||
// OpenAI
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
// Claude
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
}
|
||||
if json.Unmarshal(raw, &u) != nil {
|
||||
return
|
||||
}
|
||||
// 取较大值，避免后到的事件覆盖成 0（Claude 各事件字段不全）。
|
||||
if u.PromptTokens > p.in {
|
||||
p.in = u.PromptTokens
|
||||
}
|
||||
if u.InputTokens > p.in {
|
||||
p.in = u.InputTokens
|
||||
}
|
||||
if u.CompletionTokens > p.out {
|
||||
p.out = u.CompletionTokens
|
||||
}
|
||||
if u.OutputTokens > p.out {
|
||||
p.out = u.OutputTokens
|
||||
}
|
||||
}
|
||||
|
||||
func (p *sseParser) result() (in, out int, model string) {
|
||||
return p.in, p.out, p.model
|
||||
}
|
||||
Reference in New Issue
Block a user