Claude
@@ -0,0 +1,219 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aichat/llm"
|
||||
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
)
|
||||
|
||||
type ollamaChatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []ollamaChatMessage `json:"messages"`
|
||||
Stream bool `json:"stream"`
|
||||
Options map[string]int `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
type ollamaChatMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
Images []string `json:"images,omitempty"`
|
||||
}
|
||||
|
||||
type ollamaChatResponse struct {
|
||||
Message struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
Thinking string `json:"thinking"`
|
||||
} `json:"message"`
|
||||
Done bool `json:"done"`
|
||||
PromptEvalCount int `json:"prompt_eval_count"`
|
||||
EvalCount int `json:"eval_count"`
|
||||
DoneReason string `json:"done_reason"`
|
||||
}
|
||||
|
||||
func StreamOllamaChat(ctx context.Context, profile *llm.Profile, messages []*model.ChatCompletionMessage, promptTokens int, usage *Tracker, emit EmitFunc, onDone func(string)) error {
|
||||
requestMessages, err := buildOllamaMessages(messages)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
baseURL, err := llm.OllamaBaseURL(profile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
body, err := json.Marshal(ollamaChatRequest{
|
||||
Model: profile.Config.Model,
|
||||
Messages: requestMessages,
|
||||
Stream: true,
|
||||
Options: map[string]int{"num_predict": 4096},
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(baseURL, "/")+"/api/chat", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
data, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return fmt.Errorf("Ollama 原生接口调用失败: %s %s", resp.Status, strings.TrimSpace(string(data)))
|
||||
}
|
||||
|
||||
emit(Frame{Type: "trace", Tool: "model", Stage: "stream", Status: "running", Message: "Ollama 视觉模型已开始输出"})
|
||||
parseThinkTags := llm.ShouldParseThinkTags(profile)
|
||||
thinkParser := &Parser{}
|
||||
var full strings.Builder
|
||||
completionTokens := 0
|
||||
streamStarted := time.Now()
|
||||
peakTokensPerSecond := 0.0
|
||||
emitDelta := func(delta string) {
|
||||
if delta == "" {
|
||||
return
|
||||
}
|
||||
full.WriteString(delta)
|
||||
completionTokens += EstimateTokenCount(delta)
|
||||
usage.SetModel(promptTokens, completionTokens)
|
||||
currentSpeed := TokensPerSecond(completionTokens, streamStarted)
|
||||
if currentSpeed > peakTokensPerSecond {
|
||||
peakTokensPerSecond = currentSpeed
|
||||
}
|
||||
stats := usage.Snapshot(currentSpeed, peakTokensPerSecond)
|
||||
emit(Frame{Type: "delta", Text: delta, Stats: &stats})
|
||||
}
|
||||
emitContent := func(delta string) {
|
||||
if delta == "" {
|
||||
return
|
||||
}
|
||||
if !parseThinkTags {
|
||||
emitDelta(delta)
|
||||
return
|
||||
}
|
||||
visible, reasoning := thinkParser.Accept(delta)
|
||||
if reasoning != "" {
|
||||
emit(Frame{Type: "reasoning", Text: reasoning})
|
||||
}
|
||||
emitDelta(visible)
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024)
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
var chunk ollamaChatResponse
|
||||
if err := json.Unmarshal([]byte(line), &chunk); err != nil {
|
||||
return fmt.Errorf("解析 Ollama 流失败: %w", err)
|
||||
}
|
||||
if chunk.Message.Thinking != "" {
|
||||
emit(Frame{Type: "reasoning", Text: chunk.Message.Thinking})
|
||||
}
|
||||
emitContent(chunk.Message.Content)
|
||||
if chunk.Done {
|
||||
if chunk.PromptEvalCount > 0 || chunk.EvalCount > 0 {
|
||||
usage.SetModel(chunk.PromptEvalCount, chunk.EvalCount)
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if parseThinkTags {
|
||||
visible, reasoning := thinkParser.Flush()
|
||||
if reasoning != "" {
|
||||
emit(Frame{Type: "reasoning", Text: reasoning})
|
||||
}
|
||||
emitDelta(visible)
|
||||
}
|
||||
if onDone != nil {
|
||||
onDone(full.String())
|
||||
}
|
||||
finalStats := usage.Snapshot(TokensPerSecond(completionTokens, streamStarted), peakTokensPerSecond)
|
||||
emit(Frame{Type: "stats", Stats: &finalStats})
|
||||
emit(Frame{Type: "trace", Tool: "model", Stage: "stream", Status: "success", Message: "回答生成完成"})
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildOllamaMessages(messages []*model.ChatCompletionMessage) ([]ollamaChatMessage, error) {
|
||||
result := make([]ollamaChatMessage, 0, len(messages))
|
||||
for _, msg := range messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
role := string(msg.Role)
|
||||
if msg.Role == model.ChatMessageRoleTool {
|
||||
role = string(model.ChatMessageRoleUser)
|
||||
}
|
||||
item := ollamaChatMessage{Role: role}
|
||||
if msg.Content == nil {
|
||||
if len(msg.ToolCalls) > 0 {
|
||||
continue
|
||||
}
|
||||
result = append(result, item)
|
||||
continue
|
||||
}
|
||||
if msg.Content.StringValue != nil {
|
||||
item.Content = *msg.Content.StringValue
|
||||
if msg.Role == model.ChatMessageRoleTool {
|
||||
item.Content = "工具结果：\n" + item.Content
|
||||
}
|
||||
result = append(result, item)
|
||||
continue
|
||||
}
|
||||
for _, part := range msg.Content.ListValue {
|
||||
if part == nil {
|
||||
continue
|
||||
}
|
||||
switch part.Type {
|
||||
case model.ChatCompletionMessageContentPartTypeText:
|
||||
if part.Text != "" {
|
||||
if item.Content != "" {
|
||||
item.Content += "\n"
|
||||
}
|
||||
item.Content += part.Text
|
||||
}
|
||||
case model.ChatCompletionMessageContentPartTypeImageURL:
|
||||
if part.ImageURL == nil {
|
||||
continue
|
||||
}
|
||||
image, err := ollamaImagePayload(part.ImageURL.URL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.Images = append(item.Images, image)
|
||||
}
|
||||
}
|
||||
result = append(result, item)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func ollamaImagePayload(raw string) (string, error) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if strings.HasPrefix(strings.ToLower(raw), "data:") {
|
||||
comma := strings.Index(raw, ",")
|
||||
if comma < 0 {
|
||||
return "", errors.New("图片 base64 数据格式错误")
|
||||
}
|
||||
return strings.TrimSpace(raw[comma+1:]), nil
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
Reference in New Issue
Block a user