This commit is contained in:
renzhiyuan 2025-03-27 10:52:35 +08:00
commit c961144858
5 changed files with 118 additions and 0 deletions

24
ai/ai.go Normal file
View File

@ -0,0 +1,24 @@
package l_ai_address
import (
"context"
"gitea.cdlsxd.cn/self-tools/l_ai_brand/ai/doubao"
"strings"
)
func GetBrand(ctx context.Context, goodsName, key, model string, brandList []string) (brand string, err error) {
modelObj := doubao.NewDouBao(model, key)
text := []string{
"根据商品名称,获取对应品牌[QUESTION]" + goodsName + "[/QUESTION]",
"-只需要获取一个品牌名称",
"-如果无法获取,则返回`失败`",
}
if brandList != nil {
text = append(text, "-优先从以下品牌选取:"+strings.Join(brandList, "、"))
}
brand, err = modelObj.GetData(ctx, key, func(input string) (string, error) {
return input, nil
}, text...)
return
}

12
ai/ai_test.go Normal file
View File

@ -0,0 +1,12 @@
package l_ai_address
import (
"context"
"strings"
"testing"
)
func TestAddress(t *testing.T) {
res, err := GetBrand(context.Background(), "促销 维达Vinda抽纸 超韧130抽*24包S码 湿水不易破 卫生纸 纸巾餐巾纸 整箱", "236ba4b6-9daa-4755-b22f-2fd274cd223a", "doubao-1-5-lite-32k-250115", strings.Split("爱奇艺,QQ阅读,优酷,芒果,腾讯,QQ音乐,KFC,奈雪的茶,星巴克,Keep,盒马生鲜,COCO,E袋洗,NN加速器,PP视频,WPS,百度网盘,百度文库,必胜客,哔哩哔哩,得到,滴滴,饿了么,哈根达斯,哈罗出行,京东,酷狗音乐,酷我音乐,懒人听书,良品铺子,埋堆堆,麦当劳,猫眼电影,美团外卖,咪咕视频,咪咕文学,面包新语,南瓜电影,十点读书,搜狐视频,苏宁易购,淘票票,腾讯QQ,天猫,网易云音乐,沃尔玛,喜马拉雅,迅雷,知乎,中国电信,中国联通,中国石化,中国移动,,九阳,奔腾,炊大皇,西屋,小熊,维达Vinda", ","))
t.Log(res, err)
}

17
ai/doubao/constant.go Normal file
View File

@ -0,0 +1,17 @@
package doubao
var UrlMap = map[UrlType]string{
Text: "https://ark.cn-beijing.volces.com/api/v3/chat/completions",
Video: "https://ark.cn-beijing.volces.com/api/v3/contents/generations/tasks",
Embedding: "https://ark.cn-beijing.volces.com/api/v3/embeddings",
Token: "https://ark.cn-beijing.volces.com/api/v3/tokenization",
}
type UrlType string
const (
Text UrlType = "text"
Video UrlType = "video"
Embedding UrlType = "embedding"
Token UrlType = "token"
)

53
ai/doubao/doubao.go Normal file
View File

@ -0,0 +1,53 @@
package doubao
import (
"context"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
"github.com/volcengine/volcengine-go-sdk/volcengine"
"time"
)
type DouBao struct {
Model string
Key string
}
func NewDouBao(modelName string, key string) *DouBao {
return &DouBao{
Model: modelName,
Key: key,
}
}
func (o *DouBao) GetData(ctx context.Context, key string, respHandle func(input string) (string, error), text ...string) (string, error) {
var Message = make([]*model.ChatCompletionMessage, len(text))
client := arkruntime.NewClientWithApiKey(
key,
//arkruntime.WithBaseUrl(UrlMap[url]),
arkruntime.WithRegion("cn-beijing"),
arkruntime.WithTimeout(2*time.Minute),
arkruntime.WithRetryTimes(2),
)
for k, v := range text {
Message[k] = &model.ChatCompletionMessage{
Role: model.ChatMessageRoleSystem,
Content: &model.ChatCompletionMessageContent{
StringValue: volcengine.String(v),
},
}
}
req := model.CreateChatCompletionRequest{
Model: o.Model,
Messages: Message,
}
resp, err := client.CreateChatCompletion(ctx, req)
if err != nil {
return "", err
}
result, err := respHandle(*resp.Choices[0].Message.Content.StringValue)
return result, err
}

12
go.mod Normal file
View File

@ -0,0 +1,12 @@
module gitea.cdlsxd.cn/self-tools/l_ai_brand
go 1.23.6
require github.com/volcengine/volcengine-go-sdk v1.0.187
require (
github.com/google/uuid v1.3.0 // indirect
github.com/jmespath/go-jmespath v0.4.0 // indirect
github.com/volcengine/volc-sdk-golang v1.0.23 // indirect
gopkg.in/yaml.v2 v2.2.8 // indirect
)