74 lines
2.2 KiB
Go
74 lines
2.2 KiB
Go
package completion
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"meshtastic_mqtt_server/llm"
|
|
"meshtastic_mqtt_server/message"
|
|
|
|
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
|
)
|
|
|
|
// ChatCompleter is a function that completes a chat
|
|
type ChatCompleter func(ctx context.Context, profile *llm.Profile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error)
|
|
|
|
// CompleteChat completes a chat conversation
|
|
func CompleteChat(ctx context.Context, profile *llm.Profile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
|
|
if profile == nil || profile.Client == nil {
|
|
return model.ChatCompletionResponse{}, fmt.Errorf("llm profile or client is nil")
|
|
}
|
|
|
|
// Use context with timeout if provided
|
|
if timeout > 0 {
|
|
var cancel context.CancelFunc
|
|
ctx, cancel = context.WithTimeout(ctx, timeout)
|
|
defer cancel()
|
|
}
|
|
|
|
resp, err := profile.Client.CreateChatCompletion(ctx, req)
|
|
if err != nil {
|
|
return model.ChatCompletionResponse{}, fmt.Errorf("chat completion failed: %w", err)
|
|
}
|
|
return resp, nil
|
|
}
|
|
|
|
// CompleteText completes a text prompt using conversation messages
|
|
func CompleteText(ctx context.Context, profile *llm.Profile, messages []message.ChatMessage, maxTokens int) (string, error) {
|
|
if profile == nil || profile.Client == nil {
|
|
return "", fmt.Errorf("llm profile or client is nil")
|
|
}
|
|
|
|
arkMessages := make([]*model.ChatCompletionMessage, 0, len(messages))
|
|
for _, msg := range messages {
|
|
content := &model.ChatCompletionMessageContent{
|
|
StringValue: &msg.Content,
|
|
}
|
|
arkMessages = append(arkMessages, &model.ChatCompletionMessage{
|
|
Role: msg.Role,
|
|
Content: content,
|
|
})
|
|
}
|
|
|
|
req := model.CreateChatCompletionRequest{
|
|
Model: profile.Config.Model,
|
|
Messages: arkMessages,
|
|
MaxTokens: &maxTokens,
|
|
}
|
|
|
|
resp, err := profile.Client.CreateChatCompletion(ctx, req)
|
|
if err != nil {
|
|
return "", fmt.Errorf("text completion failed: %w", err)
|
|
}
|
|
|
|
if len(resp.Choices) == 0 {
|
|
return "", fmt.Errorf("no completion choices returned")
|
|
}
|
|
|
|
if resp.Choices[0].Message.Content != nil && resp.Choices[0].Message.Content.StringValue != nil {
|
|
return *resp.Choices[0].Message.Content.StringValue, nil
|
|
}
|
|
return "", nil
|
|
}
|