85 lines
2.4 KiB
Go
85 lines
2.4 KiB
Go
package completion
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"strings"
|
|
"time"
|
|
|
|
"aichat/llm"
|
|
"aichat/message"
|
|
"aichat/stream"
|
|
"aichat/utils"
|
|
|
|
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
|
)
|
|
|
|
type ChatCompleter func(context.Context, *llm.Profile, model.CreateChatCompletionRequest, time.Duration) (model.ChatCompletionResponse, error)
|
|
|
|
func CompleteText(ctx context.Context, profile *llm.Profile, chatMessages []message.ChatMessage, maxTokens int) (string, error) {
|
|
return CompleteTextWithTimeout(ctx, profile, chatMessages, maxTokens, time.Duration(profile.Config.Timeout)*time.Second)
|
|
}
|
|
|
|
func CompleteTextWithTimeout(ctx context.Context, profile *llm.Profile, chatMessages []message.ChatMessage, maxTokens int, timeout time.Duration) (string, error) {
|
|
messages, err := message.BuildArkMessages(chatMessages)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
completionCtx, cancel := context.WithTimeout(ctx, timeout)
|
|
defer cancel()
|
|
streamResp, err := profile.Client.CreateChatCompletionStream(completionCtx, model.CreateChatCompletionRequest{
|
|
Model: profile.Config.Model,
|
|
Messages: messages,
|
|
MaxTokens: utils.IntPtr(maxTokens),
|
|
}.WithStream(true))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer streamResp.Close()
|
|
|
|
promptTokens := stream.EstimateChatMessagesTokens(chatMessages)
|
|
completionTokens := 0
|
|
parseThinkTags := llm.ShouldParseThinkTags(profile)
|
|
thinkParser := &stream.Parser{}
|
|
var b strings.Builder
|
|
appendVisible := func(delta string) {
|
|
if delta == "" {
|
|
return
|
|
}
|
|
b.WriteString(delta)
|
|
completionTokens += stream.EstimateTokenCount(delta)
|
|
}
|
|
for {
|
|
resp, err := streamResp.Recv()
|
|
if errors.Is(err, io.EOF) {
|
|
if parseThinkTags {
|
|
visible, _ := thinkParser.Flush()
|
|
appendVisible(visible)
|
|
}
|
|
if tracker := stream.TrackerFromContext(ctx); tracker != nil {
|
|
tracker.AddTool(promptTokens, completionTokens)
|
|
}
|
|
return b.String(), nil
|
|
}
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if len(resp.Choices) > 0 {
|
|
delta := resp.Choices[0].Delta.Content
|
|
if parseThinkTags {
|
|
visible, _ := thinkParser.Accept(delta)
|
|
appendVisible(visible)
|
|
} else {
|
|
appendVisible(delta)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func CompleteChatWithTimeout(ctx context.Context, profile *llm.Profile, request model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
|
|
completionCtx, cancel := context.WithTimeout(ctx, timeout)
|
|
defer cancel()
|
|
return profile.Client.CreateChatCompletion(completionCtx, request.WithStream(false))
|
|
}
|