mirror of
https://github.com/httprunner/httprunner.git
synced 2026-05-31 21:39:41 +08:00
refactor: split ai related logic to pkg/ai
This commit is contained in:
lilong.129
63
pkg/ai/ai.go
Normal file
63
pkg/ai/ai.go
Normal file
@@ -0,0 +1,63 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
"github.com/httprunner/httprunner/v5/code"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
func NewAIService(opts ...AIServiceOption) *AIServices {
|
||||
services := &AIServices{}
|
||||
for _, option := range opts {
|
||||
option(services)
|
||||
}
|
||||
return services
|
||||
}
|
||||
|
||||
type AIServices struct {
|
||||
ICVService
|
||||
ILLMService
|
||||
}
|
||||
|
||||
type AIServiceOption func(*AIServices)
|
||||
|
||||
type CVServiceType string
|
||||
|
||||
const (
|
||||
CVServiceTypeVEDEM CVServiceType = "vedem"
|
||||
CVServiceTypeOpenCV CVServiceType = "opencv"
|
||||
)
|
||||
|
||||
func WithCVService(service CVServiceType) AIServiceOption {
|
||||
return func(opts *AIServices) {
|
||||
if service == CVServiceTypeVEDEM {
|
||||
var err error
|
||||
opts.ICVService, err = NewVEDEMImageService()
|
||||
if err != nil {
|
||||
log.Error().Err(err).Msg("init vedem image service failed")
|
||||
os.Exit(code.GetErrorCode(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type LLMServiceType string
|
||||
|
||||
const (
|
||||
LLMServiceTypeGPT4o LLMServiceType = "gpt-4o"
|
||||
LLMServiceTypeDeepSeekV3 LLMServiceType = "deepseek-v3"
|
||||
)
|
||||
|
||||
func WithLLMService(service LLMServiceType) AIServiceOption {
|
||||
return func(opts *AIServices) {
|
||||
if service == LLMServiceTypeGPT4o {
|
||||
var err error
|
||||
opts.ILLMService, err = NewGPT4oLLMService()
|
||||
if err != nil {
|
||||
log.Error().Err(err).Msg("init gpt-4o llm service failed")
|
||||
os.Exit(code.GetErrorCode(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
11
pkg/ai/ai_test.go
Normal file
11
pkg/ai/ai_test.go
Normal file
@@ -0,0 +1,11 @@
|
||||
package ai
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestOption(t *testing.T) {
|
||||
options := NewAIService(
|
||||
WithCVService(CVServiceTypeOpenCV),
|
||||
WithLLMService(LLMServiceTypeDeepSeekV3),
|
||||
)
|
||||
t.Log(options)
|
||||
}
|
||||
330
pkg/ai/cv.go
Normal file
330
pkg/ai/cv.go
Normal file
@@ -0,0 +1,330 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"image"
|
||||
"math"
|
||||
"regexp"
|
||||
|
||||
"github.com/httprunner/httprunner/v5/code"
|
||||
"github.com/httprunner/httprunner/v5/internal/builtin"
|
||||
"github.com/pkg/errors"
|
||||
)
|
||||
|
||||
type ICVService interface {
|
||||
// returns CV result including ocr texts, uploaded image url, etc
|
||||
ReadFromBuffer(imageBuf *bytes.Buffer, opts ...ScreenShotOption) (*CVResult, error)
|
||||
ReadFromPath(imagePath string, opts ...ScreenShotOption) (*CVResult, error)
|
||||
}
|
||||
|
||||
type CVResult struct {
|
||||
URL string `json:"url,omitempty"` // image uploaded url
|
||||
OCRResult OCRResults `json:"ocrResult,omitempty"` // OCR texts
|
||||
// NoLive（非直播间）
|
||||
// Shop（电商）
|
||||
// LifeService（生活服务）
|
||||
// Show（秀场）
|
||||
// Game（游戏）
|
||||
// People（多人）
|
||||
// PK（PK）
|
||||
// Media（媒体）
|
||||
// Chat（语音）
|
||||
// Event（赛事）
|
||||
LiveType string `json:"liveType,omitempty"` // 直播间类型
|
||||
LivePopularity int64 `json:"livePopularity,omitempty"` // 直播间热度
|
||||
UIResult UIResultMap `json:"uiResult,omitempty"` // 图标检测
|
||||
ClosePopupsResult *ClosePopupsResult `json:"closeResult,omitempty"` // 弹窗按钮检测
|
||||
}
|
||||
|
||||
type OCRResult struct {
|
||||
Text string `json:"text"`
|
||||
Points []PointF `json:"points"`
|
||||
}
|
||||
|
||||
type OCRResults []OCRResult
|
||||
|
||||
func (o OCRResults) ToOCRTexts() (ocrTexts OCRTexts) {
|
||||
for _, ocrResult := range o {
|
||||
rect := image.Rectangle{
|
||||
// ocrResult.Points 顺序：左上 -> 右上 -> 右下 -> 左下
|
||||
Min: image.Point{
|
||||
X: int(ocrResult.Points[0].X),
|
||||
Y: int(ocrResult.Points[0].Y),
|
||||
},
|
||||
Max: image.Point{
|
||||
X: int(ocrResult.Points[2].X),
|
||||
Y: int(ocrResult.Points[2].Y),
|
||||
},
|
||||
}
|
||||
rectStr := fmt.Sprintf("%d,%d,%d,%d",
|
||||
rect.Min.X, rect.Min.Y, rect.Max.X, rect.Max.Y)
|
||||
ocrText := OCRText{
|
||||
Text: ocrResult.Text,
|
||||
Rect: rect,
|
||||
RectStr: rectStr,
|
||||
}
|
||||
ocrTexts = append(ocrTexts, ocrText)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
type OCRText struct {
|
||||
Text string `json:"text"`
|
||||
RectStr string `json:"rect"`
|
||||
Rect image.Rectangle `json:"-"`
|
||||
}
|
||||
|
||||
func (t OCRText) Size() Size {
|
||||
return Size{
|
||||
Width: t.Rect.Dx(),
|
||||
Height: t.Rect.Dy(),
|
||||
}
|
||||
}
|
||||
|
||||
func (t OCRText) Center() PointF {
|
||||
return getRectangleCenterPoint(t.Rect)
|
||||
}
|
||||
|
||||
func getRectangleCenterPoint(rect image.Rectangle) (point PointF) {
|
||||
x, y := float64(rect.Min.X), float64(rect.Min.Y)
|
||||
width, height := float64(rect.Dx()), float64(rect.Dy())
|
||||
point = PointF{
|
||||
X: x + width*0.5,
|
||||
Y: y + height*0.5,
|
||||
}
|
||||
return point
|
||||
}
|
||||
|
||||
type OCRTexts []OCRText
|
||||
|
||||
func (t OCRTexts) texts() (texts []string) {
|
||||
for _, text := range t {
|
||||
texts = append(texts, text.Text)
|
||||
}
|
||||
return texts
|
||||
}
|
||||
|
||||
func (t OCRTexts) FilterScope(scope AbsScope) (results OCRTexts) {
|
||||
for _, ocrText := range t {
|
||||
rect := ocrText.Rect
|
||||
|
||||
// check if text in scope
|
||||
if len(scope) == 4 {
|
||||
if rect.Min.X < scope[0] ||
|
||||
rect.Min.Y < scope[1] ||
|
||||
rect.Max.X > scope[2] ||
|
||||
rect.Max.Y > scope[3] {
|
||||
// not in scope
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
results = append(results, ocrText)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// FindText returns matched text with options
|
||||
// Notice: filter scope should be specified with WithAbsScope
|
||||
func (t OCRTexts) FindText(text string, opts ...ScreenFilterOption) (result OCRText, err error) {
|
||||
options := NewScreenFilterOptions(opts...)
|
||||
|
||||
var results []OCRText
|
||||
for _, ocrText := range t.FilterScope(options.AbsScope) {
|
||||
if options.Regex {
|
||||
// regex on, check if match regex
|
||||
if !regexp.MustCompile(text).MatchString(ocrText.Text) {
|
||||
continue
|
||||
}
|
||||
} else {
|
||||
// regex off, check if match exactly
|
||||
if ocrText.Text != text {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
results = append(results, ocrText)
|
||||
|
||||
// return the first one matched exactly when index not specified
|
||||
if ocrText.Text == text && options.Index == 0 {
|
||||
return ocrText, nil
|
||||
}
|
||||
}
|
||||
|
||||
if len(results) == 0 {
|
||||
return OCRText{}, errors.Wrap(code.CVResultNotFoundError,
|
||||
fmt.Sprintf("text %s not found in %v", text, t.texts()))
|
||||
}
|
||||
|
||||
// get index
|
||||
idx := options.Index
|
||||
if idx < 0 {
|
||||
idx = len(results) + idx
|
||||
}
|
||||
|
||||
// index out of range
|
||||
if idx >= len(results) || idx < 0 {
|
||||
return OCRText{}, errors.Wrap(code.CVResultNotFoundError,
|
||||
fmt.Sprintf("text %s found %d, index %d out of range", text, len(results), idx))
|
||||
}
|
||||
|
||||
return results[idx], nil
|
||||
}
|
||||
|
||||
func (t OCRTexts) FindTexts(texts []string, opts ...ScreenFilterOption) (results OCRTexts, err error) {
|
||||
options := NewScreenFilterOptions(opts...)
|
||||
for _, text := range texts {
|
||||
ocrText, err := t.FindText(text, opts...)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
results = append(results, ocrText)
|
||||
|
||||
// found one, skip searching and return
|
||||
if options.MatchOne {
|
||||
return results, nil
|
||||
}
|
||||
}
|
||||
|
||||
if len(results) == len(texts) {
|
||||
return results, nil
|
||||
}
|
||||
return nil, errors.Wrap(code.CVResultNotFoundError,
|
||||
fmt.Sprintf("texts %s not found in %v", texts, t.texts()))
|
||||
}
|
||||
|
||||
type UIResultMap map[string]UIResults
|
||||
|
||||
// FilterUIResults filters ui icons, the former the uiTypes, the higher the priority
|
||||
func (u UIResultMap) FilterUIResults(uiTypes []string) (uiResults UIResults, err error) {
|
||||
var ok bool
|
||||
for _, uiType := range uiTypes {
|
||||
uiResults, ok = u[uiType]
|
||||
if ok && len(uiResults) != 0 {
|
||||
return
|
||||
}
|
||||
}
|
||||
err = errors.Wrap(code.CVResultNotFoundError, fmt.Sprintf("UI types %v not detected", uiTypes))
|
||||
return
|
||||
}
|
||||
|
||||
type UIResult struct {
|
||||
Box
|
||||
}
|
||||
|
||||
type Box struct {
|
||||
Point PointF `json:"point"`
|
||||
Width float64 `json:"width"`
|
||||
Height float64 `json:"height"`
|
||||
}
|
||||
|
||||
func (box Box) IsEmpty() bool {
|
||||
return builtin.IsZeroFloat64(box.Width) && builtin.IsZeroFloat64(box.Height)
|
||||
}
|
||||
|
||||
func (box Box) IsIdentical(box2 Box) bool {
|
||||
// set the coordinate precision to 1 pixel
|
||||
return box.Point.IsIdentical(box2.Point) &&
|
||||
builtin.IsZeroFloat64(math.Abs(box.Width-box2.Width)) &&
|
||||
builtin.IsZeroFloat64(math.Abs(box.Height-box2.Height))
|
||||
}
|
||||
|
||||
func (box Box) Center() PointF {
|
||||
return PointF{
|
||||
X: box.Point.X + box.Width*0.5,
|
||||
Y: box.Point.Y + box.Height*0.5,
|
||||
}
|
||||
}
|
||||
|
||||
type UIResults []UIResult
|
||||
|
||||
func (u UIResults) FilterScope(scope AbsScope) (results UIResults) {
|
||||
for _, uiResult := range u {
|
||||
rect := image.Rectangle{
|
||||
Min: image.Point{
|
||||
X: int(uiResult.Point.X),
|
||||
Y: int(uiResult.Point.Y),
|
||||
},
|
||||
Max: image.Point{
|
||||
X: int(uiResult.Point.X + uiResult.Width),
|
||||
Y: int(uiResult.Point.Y + uiResult.Height),
|
||||
},
|
||||
}
|
||||
|
||||
// check if ui result in scope
|
||||
if len(scope) == 4 {
|
||||
if rect.Min.X < scope[0] ||
|
||||
rect.Min.Y < scope[1] ||
|
||||
rect.Max.X > scope[2] ||
|
||||
rect.Max.Y > scope[3] {
|
||||
// not in scope
|
||||
continue
|
||||
}
|
||||
}
|
||||
results = append(results, uiResult)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (u UIResults) GetUIResult(opts ...ScreenFilterOption) (UIResult, error) {
|
||||
options := NewScreenFilterOptions(opts...)
|
||||
uiResults := u.FilterScope(options.AbsScope)
|
||||
if len(uiResults) == 0 {
|
||||
return UIResult{}, errors.Wrap(code.CVResultNotFoundError,
|
||||
"ui types not found in scope")
|
||||
}
|
||||
// get index
|
||||
idx := options.Index
|
||||
if idx < 0 {
|
||||
idx = len(uiResults) + idx
|
||||
}
|
||||
|
||||
// index out of range
|
||||
if idx >= len(uiResults) || idx < 0 {
|
||||
return UIResult{}, errors.Wrap(code.CVResultNotFoundError,
|
||||
fmt.Sprintf("ui types index %d out of range", idx))
|
||||
}
|
||||
return uiResults[idx], nil
|
||||
}
|
||||
|
||||
// ClosePopupsResult represents the result of recognized popup to close
|
||||
type ClosePopupsResult struct {
|
||||
Type string `json:"type"`
|
||||
PopupArea Box `json:"popupArea"`
|
||||
CloseArea Box `json:"closeArea"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
func (c ClosePopupsResult) IsEmpty() bool {
|
||||
return c.PopupArea.IsEmpty() && c.CloseArea.IsEmpty()
|
||||
}
|
||||
|
||||
type Point struct {
|
||||
X int `json:"x"` // upper left X coordinate of selected element
|
||||
Y int `json:"y"` // upper left Y coordinate of selected element
|
||||
}
|
||||
|
||||
type PointF struct {
|
||||
X float64 `json:"x"`
|
||||
Y float64 `json:"y"`
|
||||
}
|
||||
|
||||
func (p PointF) IsIdentical(p2 PointF) bool {
|
||||
// set the coordinate precision to 1 pixel
|
||||
return math.Abs(p.X-p2.X) < 1 && math.Abs(p.Y-p2.Y) < 1
|
||||
}
|
||||
|
||||
type Size struct {
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
}
|
||||
|
||||
func (s Size) IsNil() bool {
|
||||
return s.Width == 0 && s.Height == 0
|
||||
}
|
||||
|
||||
type Screen struct {
|
||||
StatusBarSize Size `json:"statusBarSize"`
|
||||
Scale float64 `json:"scale"`
|
||||
}
|
||||
257
pkg/ai/cv_vedem.go
Normal file
257
pkg/ai/cv_vedem.go
Normal file
@@ -0,0 +1,257 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/rs/zerolog/log"
|
||||
|
||||
"github.com/httprunner/httprunner/v5/code"
|
||||
"github.com/httprunner/httprunner/v5/internal/builtin"
|
||||
"github.com/httprunner/httprunner/v5/internal/json"
|
||||
)
|
||||
|
||||
var client = &http.Client{
|
||||
Timeout: time.Second * 10,
|
||||
}
|
||||
|
||||
type APIResponseImage struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Result CVResult `json:"result"`
|
||||
}
|
||||
|
||||
func NewVEDEMImageService() (*vedemCVService, error) {
|
||||
if err := checkEnv(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &vedemCVService{}, nil
|
||||
}
|
||||
|
||||
// vedemCVService implements IImageService interface
|
||||
// actions:
|
||||
//
|
||||
// ocr - get ocr texts
|
||||
// upload - get image uploaded url
|
||||
// liveType - get live type
|
||||
// popup - get popup windows
|
||||
// close - get close popup
|
||||
// ui - get ui position by type(s)
|
||||
type vedemCVService struct{}
|
||||
|
||||
func (s *vedemCVService) ReadFromPath(imagePath string, opts ...ScreenShotOption) (
|
||||
imageResult *CVResult, err error) {
|
||||
imageBuf, err := os.ReadFile(imagePath)
|
||||
if err != nil {
|
||||
err = errors.Wrap(code.CVRequestError,
|
||||
fmt.Sprintf("read image file error: %v", err))
|
||||
return
|
||||
}
|
||||
imageResult, err = s.ReadFromBuffer(bytes.NewBuffer(imageBuf), opts...)
|
||||
return
|
||||
}
|
||||
|
||||
func (s *vedemCVService) ReadFromBuffer(imageBuf *bytes.Buffer, opts ...ScreenShotOption) (
|
||||
imageResult *CVResult, err error) {
|
||||
actionOptions := NewScreenShotOptions(opts...)
|
||||
screenshotActions := actionOptions.List()
|
||||
if len(screenshotActions) == 0 {
|
||||
// skip
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
defer func() {
|
||||
elapsed := time.Since(start).Milliseconds()
|
||||
var logger *zerolog.Event
|
||||
if err != nil {
|
||||
logger = log.Error().Err(err)
|
||||
} else {
|
||||
logger = log.Debug()
|
||||
if imageResult.URL != "" {
|
||||
logger = logger.Str("url", imageResult.URL)
|
||||
}
|
||||
if imageResult.UIResult != nil {
|
||||
logger = logger.Interface("uiResult", imageResult.UIResult)
|
||||
}
|
||||
if imageResult.ClosePopupsResult != nil {
|
||||
if imageResult.ClosePopupsResult.IsEmpty() {
|
||||
// set nil to reduce unnecessary summary info
|
||||
imageResult.ClosePopupsResult = nil
|
||||
} else {
|
||||
logger = logger.Interface("closePopupsResult", imageResult.ClosePopupsResult)
|
||||
}
|
||||
}
|
||||
}
|
||||
logger = logger.Int64("elapsed(ms)", elapsed)
|
||||
logger.Msg("get image data by veDEM")
|
||||
}()
|
||||
|
||||
bodyBuf := &bytes.Buffer{}
|
||||
bodyWriter := multipart.NewWriter(bodyBuf)
|
||||
for _, action := range screenshotActions {
|
||||
bodyWriter.WriteField("actions", action)
|
||||
}
|
||||
for _, uiType := range actionOptions.ScreenShotWithUITypes {
|
||||
bodyWriter.WriteField("uiTypes", uiType)
|
||||
}
|
||||
|
||||
// 使用高精度集群
|
||||
bodyWriter.WriteField("ocrCluster", "highPrecision")
|
||||
|
||||
if actionOptions.ScreenShotWithOCRCluster != "" {
|
||||
bodyWriter.WriteField("ocrCluster", actionOptions.ScreenShotWithOCRCluster)
|
||||
}
|
||||
|
||||
bodyWriter.WriteField("timeout", fmt.Sprintf("%v", 10))
|
||||
|
||||
formWriter, err := bodyWriter.CreateFormFile("image", "screenshot.png")
|
||||
if err != nil {
|
||||
err = errors.Wrap(code.CVRequestError,
|
||||
fmt.Sprintf("create form file error: %v", err))
|
||||
return
|
||||
}
|
||||
|
||||
size, err := formWriter.Write(imageBuf.Bytes())
|
||||
if err != nil {
|
||||
err = errors.Wrap(code.CVRequestError,
|
||||
fmt.Sprintf("write form error: %v", err))
|
||||
return
|
||||
}
|
||||
|
||||
err = bodyWriter.Close()
|
||||
if err != nil {
|
||||
err = errors.Wrap(code.CVRequestError,
|
||||
fmt.Sprintf("close body writer error: %v", err))
|
||||
return
|
||||
}
|
||||
var req *http.Request
|
||||
var resp *http.Response
|
||||
// retry 3 times
|
||||
for i := 1; i <= 3; i++ {
|
||||
copiedBodyBuf := &bytes.Buffer{}
|
||||
if _, err := copiedBodyBuf.Write(bodyBuf.Bytes()); err != nil {
|
||||
log.Error().Err(err).Msg("copy screenshot buffer failed")
|
||||
continue
|
||||
}
|
||||
|
||||
req, err = http.NewRequest("POST", os.Getenv("VEDEM_IMAGE_URL"), copiedBodyBuf)
|
||||
if err != nil {
|
||||
err = errors.Wrap(code.CVRequestError,
|
||||
fmt.Sprintf("construct request error: %v", err))
|
||||
return
|
||||
}
|
||||
|
||||
// ppe env
|
||||
// req.Header.Add("x-tt-env", "ppe_vedem_algorithm")
|
||||
// req.Header.Add("x-use-ppe", "1")
|
||||
|
||||
signToken := "UNSIGNED-PAYLOAD"
|
||||
token := builtin.Sign("auth-v2", os.Getenv("VEDEM_IMAGE_AK"), os.Getenv("VEDEM_IMAGE_SK"), []byte(signToken))
|
||||
|
||||
req.Header.Add("Agw-Auth", token)
|
||||
req.Header.Add("Agw-Auth-Content", signToken)
|
||||
req.Header.Add("Content-Type", bodyWriter.FormDataContentType())
|
||||
|
||||
start := time.Now()
|
||||
resp, err = client.Do(req)
|
||||
elapsed := time.Since(start)
|
||||
if err != nil {
|
||||
log.Error().Err(err).
|
||||
Int("imageBufSize", size).
|
||||
Msgf("request veDEM OCR service error, retry %d", i)
|
||||
continue
|
||||
}
|
||||
|
||||
logID := getLogID(resp.Header)
|
||||
statusCode := resp.StatusCode
|
||||
if statusCode != http.StatusOK {
|
||||
log.Error().
|
||||
Str("X-TT-LOGID", logID).
|
||||
Int("imageBufSize", size).
|
||||
Int("statusCode", statusCode).
|
||||
Msgf("request veDEM OCR service failed, retry %d", i)
|
||||
time.Sleep(1 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
log.Debug().
|
||||
Str("X-TT-LOGID", logID).
|
||||
Int("image_bytes", size).
|
||||
Int64("elapsed(ms)", elapsed.Milliseconds()).
|
||||
Msg("request OCR service success")
|
||||
break
|
||||
}
|
||||
if resp == nil {
|
||||
err = code.CVServiceConnectionError
|
||||
return
|
||||
}
|
||||
|
||||
defer resp.Body.Close()
|
||||
|
||||
results, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
err = errors.Wrap(code.CVResponseError,
|
||||
fmt.Sprintf("read response body error: %v", err))
|
||||
return
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
err = errors.Wrap(code.CVResponseError,
|
||||
fmt.Sprintf("unexpected response status code: %d, results: %v",
|
||||
resp.StatusCode, string(results)))
|
||||
return
|
||||
}
|
||||
|
||||
var imageResponse APIResponseImage
|
||||
err = json.Unmarshal(results, &imageResponse)
|
||||
if err != nil {
|
||||
err = errors.Wrap(code.CVResponseError,
|
||||
fmt.Sprintf("json unmarshal veDEM image response body error, response=%s", string(results)))
|
||||
return
|
||||
}
|
||||
|
||||
if imageResponse.Code != 0 {
|
||||
err = errors.Wrap(code.CVResponseError,
|
||||
fmt.Sprintf("unexpected response data code: %d, message: %s",
|
||||
imageResponse.Code, imageResponse.Message))
|
||||
return
|
||||
}
|
||||
|
||||
imageResult = &imageResponse.Result
|
||||
return imageResult, nil
|
||||
}
|
||||
|
||||
func checkEnv() error {
|
||||
vedemImageURL := os.Getenv("VEDEM_IMAGE_URL")
|
||||
if vedemImageURL == "" {
|
||||
return errors.Wrap(code.CVEnvMissedError, "VEDEM_IMAGE_URL missed")
|
||||
}
|
||||
log.Info().Str("VEDEM_IMAGE_URL", vedemImageURL).Msg("get env")
|
||||
if os.Getenv("VEDEM_IMAGE_AK") == "" {
|
||||
return errors.Wrap(code.CVEnvMissedError, "VEDEM_IMAGE_AK missed")
|
||||
}
|
||||
if os.Getenv("VEDEM_IMAGE_SK") == "" {
|
||||
return errors.Wrap(code.CVEnvMissedError, "VEDEM_IMAGE_SK missed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func getLogID(header http.Header) string {
|
||||
if len(header) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
logID, ok := header["X-Tt-Logid"]
|
||||
if !ok || len(logID) == 0 {
|
||||
return ""
|
||||
}
|
||||
return logID[0]
|
||||
}
|
||||
41
pkg/ai/cv_vedem_test.go
Normal file
41
pkg/ai/cv_vedem_test.go
Normal file
@@ -0,0 +1,41 @@
|
||||
//go:build localtest
|
||||
|
||||
package ai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"os"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGetImageFromBuffer(t *testing.T) {
|
||||
imagePath := "/Users/debugtalk/Downloads/s1.png"
|
||||
file, err := os.ReadFile(imagePath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
buf := new(bytes.Buffer)
|
||||
buf.Read(file)
|
||||
|
||||
service := NewAIService(
|
||||
WithCVService(CVServiceTypeVEDEM),
|
||||
)
|
||||
cvResult, err := service.ReadFromBuffer(buf)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fmt.Println(fmt.Sprintf("cvResult: %v", cvResult))
|
||||
}
|
||||
|
||||
func TestGetImageFromPath(t *testing.T) {
|
||||
imagePath := "/Users/debugtalk/Downloads/s1.png"
|
||||
service := NewAIService(
|
||||
WithCVService(CVServiceTypeVEDEM),
|
||||
)
|
||||
cvResult, err := service.ReadFromPath(imagePath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fmt.Println(fmt.Sprintf("cvResult: %v", cvResult))
|
||||
}
|
||||
20
pkg/ai/llm.go
Normal file
20
pkg/ai/llm.go
Normal file
@@ -0,0 +1,20 @@
|
||||
package ai
|
||||
|
||||
import "context"
|
||||
|
||||
type ILLMService interface {
|
||||
Call(ctx context.Context, prompt string) (string, error)
|
||||
}
|
||||
|
||||
func NewGPT4oLLMService() (*openaiLLMService, error) {
|
||||
if err := checkEnv(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &openaiLLMService{}, nil
|
||||
}
|
||||
|
||||
type openaiLLMService struct{}
|
||||
|
||||
func (s openaiLLMService) Call(ctx context.Context, prompt string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
213
pkg/ai/screen.go
Normal file
213
pkg/ai/screen.go
Normal file
@@ -0,0 +1,213 @@
|
||||
package ai
|
||||
|
||||
func NewScreenShotOptions(opts ...ScreenShotOption) *ScreenShotOptions {
|
||||
options := &ScreenShotOptions{}
|
||||
for _, option := range opts {
|
||||
option(options)
|
||||
}
|
||||
return options
|
||||
}
|
||||
|
||||
type ScreenShotOptions struct {
|
||||
ScreenShotWithOCR bool `json:"screenshot_with_ocr,omitempty" yaml:"screenshot_with_ocr,omitempty"`
|
||||
ScreenShotWithUpload bool `json:"screenshot_with_upload,omitempty" yaml:"screenshot_with_upload,omitempty"`
|
||||
ScreenShotWithLiveType bool `json:"screenshot_with_live_type,omitempty" yaml:"screenshot_with_live_type,omitempty"`
|
||||
ScreenShotWithLivePopularity bool `json:"screenshot_with_live_popularity,omitempty" yaml:"screenshot_with_live_popularity,omitempty"`
|
||||
ScreenShotWithUITypes []string `json:"screenshot_with_ui_types,omitempty" yaml:"screenshot_with_ui_types,omitempty"`
|
||||
ScreenShotWithClosePopups bool `json:"screenshot_with_close_popups,omitempty" yaml:"screenshot_with_close_popups,omitempty"`
|
||||
ScreenShotWithOCRCluster string `json:"screenshot_with_ocr_cluster,omitempty" yaml:"screenshot_with_ocr_cluster,omitempty"`
|
||||
ScreenShotFileName string `json:"screenshot_file_name,omitempty" yaml:"screenshot_file_name,omitempty"`
|
||||
}
|
||||
|
||||
func (o *ScreenShotOptions) Options() []ScreenShotOption {
|
||||
options := make([]ScreenShotOption, 0)
|
||||
if o == nil {
|
||||
return options
|
||||
}
|
||||
|
||||
// screenshot options
|
||||
if o.ScreenShotWithOCR {
|
||||
options = append(options, WithScreenShotOCR(true))
|
||||
}
|
||||
if o.ScreenShotWithUpload {
|
||||
options = append(options, WithScreenShotUpload(true))
|
||||
}
|
||||
if o.ScreenShotWithLiveType {
|
||||
options = append(options, WithScreenShotLiveType(true))
|
||||
}
|
||||
if o.ScreenShotWithLivePopularity {
|
||||
options = append(options, WithScreenShotLivePopularity(true))
|
||||
}
|
||||
if len(o.ScreenShotWithUITypes) > 0 {
|
||||
options = append(options, WithScreenShotUITypes(o.ScreenShotWithUITypes...))
|
||||
}
|
||||
if o.ScreenShotWithClosePopups {
|
||||
options = append(options, WithScreenShotClosePopups(true))
|
||||
}
|
||||
if o.ScreenShotWithOCRCluster != "" {
|
||||
options = append(options, WithScreenOCRCluster(o.ScreenShotWithOCRCluster))
|
||||
}
|
||||
if o.ScreenShotFileName != "" {
|
||||
options = append(options, WithScreenShotFileName(o.ScreenShotFileName))
|
||||
}
|
||||
|
||||
return options
|
||||
}
|
||||
|
||||
func (o *ScreenShotOptions) List() []string {
|
||||
options := []string{}
|
||||
if o.ScreenShotWithUpload {
|
||||
options = append(options, "upload")
|
||||
}
|
||||
if o.ScreenShotWithOCR {
|
||||
options = append(options, "ocr")
|
||||
}
|
||||
if o.ScreenShotWithLiveType {
|
||||
options = append(options, "liveType")
|
||||
}
|
||||
if o.ScreenShotWithLivePopularity {
|
||||
options = append(options, "livePopularity")
|
||||
}
|
||||
// UI detection
|
||||
if len(o.ScreenShotWithUITypes) > 0 {
|
||||
options = append(options, "ui")
|
||||
}
|
||||
if o.ScreenShotWithClosePopups {
|
||||
options = append(options, "close")
|
||||
}
|
||||
return options
|
||||
}
|
||||
|
||||
type ScreenShotOption func(o *ScreenShotOptions)
|
||||
|
||||
func WithScreenShotOCR(ocrOn bool) ScreenShotOption {
|
||||
return func(o *ScreenShotOptions) {
|
||||
o.ScreenShotWithOCR = ocrOn
|
||||
}
|
||||
}
|
||||
|
||||
func WithScreenShotUpload(uploadOn bool) ScreenShotOption {
|
||||
return func(o *ScreenShotOptions) {
|
||||
o.ScreenShotWithUpload = uploadOn
|
||||
}
|
||||
}
|
||||
|
||||
func WithScreenShotLiveType(liveTypeOn bool) ScreenShotOption {
|
||||
return func(o *ScreenShotOptions) {
|
||||
o.ScreenShotWithLiveType = liveTypeOn
|
||||
}
|
||||
}
|
||||
|
||||
func WithScreenShotLivePopularity(livePopularityOn bool) ScreenShotOption {
|
||||
return func(o *ScreenShotOptions) {
|
||||
o.ScreenShotWithLivePopularity = livePopularityOn
|
||||
}
|
||||
}
|
||||
|
||||
func WithScreenShotUITypes(uiTypes ...string) ScreenShotOption {
|
||||
return func(o *ScreenShotOptions) {
|
||||
o.ScreenShotWithUITypes = uiTypes
|
||||
}
|
||||
}
|
||||
|
||||
func WithScreenShotClosePopups(closeOn bool) ScreenShotOption {
|
||||
return func(o *ScreenShotOptions) {
|
||||
o.ScreenShotWithClosePopups = closeOn
|
||||
}
|
||||
}
|
||||
|
||||
func WithScreenOCRCluster(ocrCluster string) ScreenShotOption {
|
||||
return func(o *ScreenShotOptions) {
|
||||
o.ScreenShotWithOCRCluster = ocrCluster
|
||||
}
|
||||
}
|
||||
|
||||
func WithScreenShotFileName(fileName string) ScreenShotOption {
|
||||
return func(o *ScreenShotOptions) {
|
||||
o.ScreenShotFileName = fileName
|
||||
}
|
||||
}
|
||||
|
||||
// (x1, y1) is the top left corner, (x2, y2) is the bottom right corner
|
||||
// [x1, y1, x2, y2] in percentage of the screen
|
||||
type Scope []float64
|
||||
|
||||
func (s Scope) ToAbs(windowSize Size) AbsScope {
|
||||
x1, y1, x2, y2 := s[0], s[1], s[2], s[3]
|
||||
// convert relative scope to absolute scope
|
||||
absX1 := int(x1 * float64(windowSize.Width))
|
||||
absY1 := int(y1 * float64(windowSize.Height))
|
||||
absX2 := int(x2 * float64(windowSize.Width))
|
||||
absY2 := int(y2 * float64(windowSize.Height))
|
||||
return AbsScope{absX1, absY1, absX2, absY2}
|
||||
}
|
||||
|
||||
// [x1, y1, x2, y2] in absolute pixels
|
||||
type AbsScope []int
|
||||
|
||||
func (s AbsScope) Option() ScreenFilterOption {
|
||||
return WithAbsScope(s[0], s[1], s[2], s[3])
|
||||
}
|
||||
|
||||
func NewScreenFilterOptions(opts ...ScreenFilterOption) *ScreenFilterOptions {
|
||||
options := &ScreenFilterOptions{}
|
||||
for _, option := range opts {
|
||||
option(options)
|
||||
}
|
||||
return options
|
||||
}
|
||||
|
||||
type ScreenFilterOptions struct {
|
||||
// scope related
|
||||
Scope Scope `json:"scope,omitempty" yaml:"scope,omitempty"`
|
||||
AbsScope AbsScope `json:"abs_scope,omitempty" yaml:"abs_scope,omitempty"`
|
||||
|
||||
Regex bool `json:"regex,omitempty" yaml:"regex,omitempty"` // use regex to match text
|
||||
Offset []int `json:"offset,omitempty" yaml:"offset,omitempty"` // used to tap offset of point
|
||||
OffsetRandomRange []int `json:"offset_random_range,omitempty" yaml:"offset_random_range,omitempty"` // set random range [min, max] for tap/swipe points
|
||||
Index int `json:"index,omitempty" yaml:"index,omitempty"` // index of the target element
|
||||
MatchOne bool `json:"match_one,omitempty" yaml:"match_one,omitempty"` // match one of the targets if existed
|
||||
}
|
||||
|
||||
type ScreenFilterOption func(o *ScreenFilterOptions)
|
||||
|
||||
// WithScope inputs area of [(x1,y1), (x2,y2)]
|
||||
// x1, y1, x2, y2 are all in [0, 1], which means the relative position of the screen
|
||||
func WithScope(x1, y1, x2, y2 float64) ScreenFilterOption {
|
||||
return func(o *ScreenFilterOptions) {
|
||||
o.Scope = Scope{x1, y1, x2, y2}
|
||||
}
|
||||
}
|
||||
|
||||
// WithAbsScope inputs area of [(x1,y1), (x2,y2)]
|
||||
// x1, y1, x2, y2 are all absolute position of the screen
|
||||
func WithAbsScope(x1, y1, x2, y2 int) ScreenFilterOption {
|
||||
return func(o *ScreenFilterOptions) {
|
||||
o.AbsScope = AbsScope{x1, y1, x2, y2}
|
||||
}
|
||||
}
|
||||
|
||||
// tap [x, y] with offset [offsetX, offsetY]
|
||||
func WithTapOffset(offsetX, offsetY int) ScreenFilterOption {
|
||||
return func(o *ScreenFilterOptions) {
|
||||
o.Offset = []int{offsetX, offsetY}
|
||||
}
|
||||
}
|
||||
|
||||
func WithRegex(regex bool) ScreenFilterOption {
|
||||
return func(o *ScreenFilterOptions) {
|
||||
o.Regex = regex
|
||||
}
|
||||
}
|
||||
|
||||
func WithMatchOne(matchOne bool) ScreenFilterOption {
|
||||
return func(o *ScreenFilterOptions) {
|
||||
o.MatchOne = matchOne
|
||||
}
|
||||
}
|
||||
|
||||
func WithIndex(index int) ScreenFilterOption {
|
||||
return func(o *ScreenFilterOptions) {
|
||||
o.Index = index
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user