更新若干功能

service/ai/openai.go

@@ -37,6 +37,10 @@ type chatDelta struct {
 			Content string `json:"content"`
 		} `json:"delta"`
 	} `json:"choices"`
+	Usage *struct {
+		PromptTokens     int64 `json:"prompt_tokens"`
+		CompletionTokens int64 `json:"completion_tokens"`
+	} `json:"usage"`
 	Error *struct {
 		Message string `json:"message"`
 	} `json:"error"`
@@ -73,6 +77,11 @@ func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (
 		defer resp.Body.Close()
 		sc := bufio.NewScanner(resp.Body)
 		sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
+		var promptChars, completionChars int
+		for _, m := range messages {
+			promptChars += len(m.Content)
+		}
+		var lastUsage *Usage
 		for sc.Scan() {
 			line := strings.TrimSpace(sc.Text())
 			if !strings.HasPrefix(line, "data:") {
@@ -90,7 +99,14 @@ func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (
 				out <- Chunk{Err: errors.New(truncate(d.Error.Message, 300))}
 				return
 			}
+			if d.Usage != nil {
+				lastUsage = &Usage{
+					PromptTokens:     d.Usage.PromptTokens,
+					CompletionTokens: d.Usage.CompletionTokens,
+				}
+			}
 			if len(d.Choices) > 0 && d.Choices[0].Delta.Content != "" {
+				completionChars += len(d.Choices[0].Delta.Content)
 				select {
 				case out <- Chunk{Content: d.Choices[0].Delta.Content}:
 				case <-ctx.Done():
@@ -100,6 +116,19 @@ func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (
 		}
 		if e := sc.Err(); e != nil && ctx.Err() == nil {
 			out <- Chunk{Err: errors.New("AI_STREAM_INTERRUPTED")}
+			return
+		}
+		if lastUsage == nil {
+			// 粗估：约 4 字符 ≈ 1 token
+			lastUsage = &Usage{
+				PromptTokens:     int64((promptChars + 3) / 4),
+				CompletionTokens: int64((completionChars + 3) / 4),
+				Estimated:        true,
+			}
+		}
+		select {
+		case out <- Chunk{Usage: lastUsage}:
+		default:
 		}
 	}()
 	return out, nil

service/ai/provider.go

@@ -11,9 +11,18 @@ type Message struct {
 }
 
 // Chunk 是流式返回的一段增量内容；Err 非空表示流异常中止。
+// Usage 在流结束时可选附带 token 用量（若服务商提供）。
 type Chunk struct {
 	Content string
 	Err     error
+	Usage   *Usage
+}
+
+// Usage 一次调用的 token 用量。
+type Usage struct {
+	PromptTokens     int64
+	CompletionTokens int64
+	Estimated        bool
 }
 
 // Provider 是 AI 服务商的统一抽象。