mirror of
https://github.com/httprunner/httprunner.git
synced 2026-09-05 15:37:35 +08:00
refactor: enhance screenshot handling by introducing base64 encoding and updating related methods
This commit is contained in:
lilong.129
+26
-49
@@ -10,7 +10,6 @@ import (
|
||||
"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"
|
||||
"github.com/httprunner/httprunner/v5/uixt/ai"
|
||||
"github.com/httprunner/httprunner/v5/uixt/option"
|
||||
@@ -143,12 +142,11 @@ func (dExt *XTDriver) StartToGoal(ctx context.Context, prompt string, opts ...op
|
||||
func (dExt *XTDriver) AIAction(ctx context.Context, prompt string, opts ...option.ActionOption) (*AIExecutionResult, error) {
|
||||
log.Info().Str("prompt", prompt).Msg("performing AI action")
|
||||
|
||||
// Step 1: Take screenshot and measure time
|
||||
screenshotStartTime := time.Now()
|
||||
screenResult, err := dExt.createScreenshotWithSession(
|
||||
option.WithScreenShotFileName(builtin.GenNameWithTimestamp("%d_screenshot")),
|
||||
// Step 1: Take screenshot and convert to base64
|
||||
screenResult, err := dExt.GetScreenResult(
|
||||
option.WithScreenShotFileName("ai_action"),
|
||||
option.WithScreenShotBase64(true),
|
||||
)
|
||||
screenshotElapsed := time.Since(screenshotStartTime).Milliseconds()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -160,7 +158,7 @@ func (dExt *XTDriver) AIAction(ctx context.Context, prompt string, opts ...optio
|
||||
aiExecutionResult := &AIExecutionResult{
|
||||
Type: "action",
|
||||
ModelCallElapsed: modelCallElapsed,
|
||||
ScreenshotElapsed: screenshotElapsed,
|
||||
ScreenshotElapsed: screenResult.Elapsed,
|
||||
ImagePath: screenResult.ImagePath,
|
||||
Resolution: &screenResult.Resolution,
|
||||
PlanningResult: &planningResult.PlanningResult,
|
||||
@@ -193,13 +191,11 @@ func (dExt *XTDriver) PlanNextAction(ctx context.Context, prompt string, opts ..
|
||||
options := option.NewActionOptions(opts...)
|
||||
resetHistory := options.ResetHistory
|
||||
|
||||
// Step 1: Take screenshot
|
||||
screenshotStartTime := time.Now()
|
||||
// Use GetScreenResult to handle screenshot capture, save, and session tracking
|
||||
screenResult, err := dExt.createScreenshotWithSession(
|
||||
option.WithScreenShotFileName(builtin.GenNameWithTimestamp("%d_screenshot")),
|
||||
// Step 1: Take screenshot and convert to base64
|
||||
screenResult, err := dExt.GetScreenResult(
|
||||
option.WithScreenShotFileName("ai_planning"),
|
||||
option.WithScreenShotBase64(true),
|
||||
)
|
||||
screenshotElapsed := time.Since(screenshotStartTime).Milliseconds()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -208,12 +204,6 @@ func (dExt *XTDriver) PlanNextAction(ctx context.Context, prompt string, opts ..
|
||||
// The planning screenshot is already stored in planningResult.ScreenResult
|
||||
dExt.GetSession().GetData(true) // reset session data to exclude planning screenshot from sub-actions
|
||||
|
||||
// get screen shot buffer base64 and size
|
||||
screenShotBase64, size, err := dExt.GetScreenshotBase64WithSize()
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(code.DeviceGetInfoError, err.Error())
|
||||
}
|
||||
|
||||
// Step 2: Call model
|
||||
modelCallStartTime := time.Now()
|
||||
planningOpts := &ai.PlanningOptions{
|
||||
@@ -224,12 +214,12 @@ func (dExt *XTDriver) PlanNextAction(ctx context.Context, prompt string, opts ..
|
||||
{
|
||||
Type: schema.ChatMessagePartTypeImageURL,
|
||||
ImageURL: &schema.ChatMessageImageURL{
|
||||
URL: screenShotBase64,
|
||||
URL: screenResult.Base64,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Size: size,
|
||||
Size: screenResult.Resolution,
|
||||
ResetHistory: resetHistory,
|
||||
}
|
||||
|
||||
@@ -250,7 +240,7 @@ func (dExt *XTDriver) PlanNextAction(ctx context.Context, prompt string, opts ..
|
||||
planningResult := &PlanningExecutionResult{
|
||||
PlanningResult: *result, // Inherit all fields from ai.PlanningResult
|
||||
// Planning process timing and metadata
|
||||
ScreenshotElapsed: screenshotElapsed,
|
||||
ScreenshotElapsed: screenResult.Elapsed,
|
||||
ImagePath: screenResult.ImagePath,
|
||||
Resolution: &screenResult.Resolution,
|
||||
ScreenResult: screenResult,
|
||||
@@ -374,17 +364,11 @@ func (dExt *XTDriver) AIQuery(text string, opts ...option.ActionOption) (*AIExec
|
||||
return nil, errors.New("LLM service is not initialized")
|
||||
}
|
||||
|
||||
// Step 1: Take screenshot and measure time
|
||||
screenshotStartTime := time.Now()
|
||||
screenResult, err := dExt.createScreenshotWithSession(
|
||||
option.WithScreenShotFileName(builtin.GenNameWithTimestamp("%d_screenshot")),
|
||||
// Step 1: Take screenshot and convert to base64
|
||||
screenResult, err := dExt.GetScreenResult(
|
||||
option.WithScreenShotFileName("ai_query"),
|
||||
option.WithScreenShotBase64(true),
|
||||
)
|
||||
screenshotElapsed := time.Since(screenshotStartTime).Milliseconds()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
screenShotBase64, size, err := dExt.GetScreenshotBase64WithSize()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -398,8 +382,8 @@ func (dExt *XTDriver) AIQuery(text string, opts ...option.ActionOption) (*AIExec
|
||||
// execute query
|
||||
queryOpts := &ai.QueryOptions{
|
||||
Query: text,
|
||||
Screenshot: screenShotBase64,
|
||||
Size: size,
|
||||
Screenshot: screenResult.Base64,
|
||||
Size: screenResult.Resolution,
|
||||
OutputSchema: actionOptions.OutputSchema,
|
||||
}
|
||||
result, err := dExt.LLMService.Query(context.Background(), queryOpts)
|
||||
@@ -412,7 +396,7 @@ func (dExt *XTDriver) AIQuery(text string, opts ...option.ActionOption) (*AIExec
|
||||
aiResult := &AIExecutionResult{
|
||||
Type: "query",
|
||||
ModelCallElapsed: modelCallElapsed, // model call timing
|
||||
ScreenshotElapsed: screenshotElapsed, // screenshot timing
|
||||
ScreenshotElapsed: screenResult.Elapsed, // screenshot timing
|
||||
ImagePath: screenResult.ImagePath, // screenshot path
|
||||
Resolution: &screenResult.Resolution, // screen resolution
|
||||
QueryResult: result, // query-specific result
|
||||
@@ -426,35 +410,28 @@ func (dExt *XTDriver) AIAssert(assertion string, opts ...option.ActionOption) (*
|
||||
return nil, errors.New("LLM service is not initialized")
|
||||
}
|
||||
|
||||
// Step 1: Take screenshot and measure time
|
||||
screenshotStartTime := time.Now()
|
||||
screenResult, err := dExt.createScreenshotWithSession(
|
||||
option.WithScreenShotFileName(builtin.GenNameWithTimestamp("%d_screenshot")),
|
||||
// Step 1: Take screenshot and convert to base64
|
||||
screenResult, err := dExt.GetScreenResult(
|
||||
option.WithScreenShotFileName("ai_assert"),
|
||||
option.WithScreenShotBase64(true),
|
||||
)
|
||||
screenshotElapsed := time.Since(screenshotStartTime).Milliseconds()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
assertResult := &AIExecutionResult{
|
||||
Type: "assert",
|
||||
ScreenshotElapsed: screenshotElapsed,
|
||||
ScreenshotElapsed: screenResult.Elapsed,
|
||||
ImagePath: screenResult.ImagePath,
|
||||
Resolution: &screenResult.Resolution,
|
||||
}
|
||||
|
||||
screenShotBase64, size, err := dExt.GetScreenshotBase64WithSize()
|
||||
if err != nil {
|
||||
assertResult.Error = err.Error()
|
||||
return assertResult, err
|
||||
}
|
||||
|
||||
// Step 2: Call model and measure time
|
||||
modelCallStartTime := time.Now()
|
||||
assertOpts := &ai.AssertOptions{
|
||||
Assertion: assertion,
|
||||
Screenshot: screenShotBase64,
|
||||
Size: size,
|
||||
Screenshot: screenResult.Base64,
|
||||
Size: screenResult.Resolution,
|
||||
}
|
||||
result, err := dExt.LLMService.Assert(context.Background(), assertOpts)
|
||||
assertResult.ModelCallElapsed = time.Since(modelCallStartTime).Milliseconds()
|
||||
|
||||
Reference in New Issue
Block a user