76 lines
1.8 KiB
Go
76 lines
1.8 KiB
Go
package knowledge_base
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
)
|
|
|
|
type openAIChunk struct {
|
|
Choices []struct {
|
|
Delta *Delta `json:"delta"`
|
|
Message *Message `json:"message"`
|
|
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
|
|
}
|
|
|
|
type Message struct {
|
|
Role string `json:"role"` // 角色
|
|
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
|
|
}
|
|
|
|
func ParseOpenAIHTTPData(body string) (message *Message, done bool, err error) {
|
|
data := strings.TrimSpace(body)
|
|
if data == "" {
|
|
return nil, false, nil
|
|
}
|
|
|
|
var resp openAIChunk
|
|
if err := json.Unmarshal([]byte(data), &resp); err != nil {
|
|
return nil, false, fmt.Errorf("unmarshal openai stream chunk failed: %w", err)
|
|
}
|
|
|
|
for _, c := range resp.Choices {
|
|
if c.Message != nil {
|
|
return c.Message, true, nil // 只输出第一个message
|
|
}
|
|
}
|
|
|
|
return nil, false, nil
|
|
}
|