61 lines
1.4 KiB
Go
61 lines
1.4 KiB
Go
package utils_vllm
|
|
|
|
import (
|
|
"ai_scheduler/internal/config"
|
|
"context"
|
|
|
|
"github.com/cloudwego/eino-ext/components/model/openai"
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
type Client struct {
|
|
model *openai.ChatModel
|
|
config *config.Config
|
|
}
|
|
|
|
func NewClient(config *config.Config) (*Client, func(), error) {
|
|
m, err := openai.NewChatModel(context.Background(), &openai.ChatModelConfig{
|
|
BaseURL: config.Vllm.BaseURL,
|
|
Model: config.Vllm.VlModel,
|
|
Timeout: config.Vllm.Timeout,
|
|
})
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
c := &Client{model: m, config: config}
|
|
cleanup := func() {}
|
|
return c, cleanup, nil
|
|
}
|
|
|
|
func (c *Client) Chat(ctx context.Context, msgs []*schema.Message) (*schema.Message, error) {
|
|
return c.model.Generate(ctx, msgs)
|
|
}
|
|
|
|
func (c *Client) RecognizeWithImg(ctx context.Context, systemPrompt, userPrompt string, imgURLs []string) (*schema.Message, error) {
|
|
in := []*schema.Message{
|
|
{
|
|
Role: schema.System,
|
|
Content: systemPrompt,
|
|
},
|
|
{
|
|
Role: schema.User,
|
|
},
|
|
}
|
|
parts := []schema.MessageInputPart{
|
|
{Type: schema.ChatMessagePartTypeText, Text: userPrompt},
|
|
}
|
|
for i := range imgURLs {
|
|
u := imgURLs[i]
|
|
parts = append(parts, schema.MessageInputPart{
|
|
Type: schema.ChatMessagePartTypeImageURL,
|
|
Image: &schema.MessageInputImage{
|
|
MessagePartCommon: schema.MessagePartCommon{URL: &u},
|
|
Detail: schema.ImageURLDetailHigh,
|
|
},
|
|
})
|
|
}
|
|
|
|
in[1].UserInputMultiContent = parts
|
|
return c.model.Generate(ctx, in)
|
|
}
|