ai_scheduler/internal/domain/tools/common/knowledge_base/parse.go

49 lines
1.1 KiB
Go

package knowledge_base
import (
"encoding/json"
"fmt"
"strings"
)
type openAIChunk struct {
Choices []struct {
Delta *Delta `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
}
type Delta struct {
ReasoningContent string `json:"reasoning_content"` // 推理内容
Content string `json:"content"` // 内容
XRagStatus string `json:"x_rag_status"` // rag命中状态 hit|miss
}
func ParseOpenAIStreamData(dataLine string) (delta *Delta, done bool, err error) {
data := strings.TrimSpace(strings.TrimPrefix(dataLine, "data:"))
if data == "" {
return nil, false, nil
}
data = strings.TrimSpace(data)
if data == "" {
return nil, false, nil
}
if data == "[DONE]" {
return nil, true, nil
}
var chunk openAIChunk
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
return nil, false, fmt.Errorf("unmarshal openai stream chunk failed: %w", err)
}
for _, c := range chunk.Choices {
if c.Delta != nil {
return c.Delta, false, nil // 只输出第一个delta
}
}
return nil, false, nil
}