feat: 实现 AIQuery 功能并支持 OutputSchema

- 新增 AIQuery 方法到 StepMobile，支持使用自然语言从屏幕中提取信息
- 实现 AIQuery 在 driver_ext_ai.go 中的完整功能，包括屏幕截图和 LLM 查询
- 添加 OutputSchema 支持，允许用户定义自定义输出格式进行结构化查询
- 新增 ToolAIQuery MCP 工具，完整集成到 MCP 服务器中
- 在 ActionOptions 中添加 OutputSchema 字段和 WithOutputSchema 选项函数
- 添加 ACTION_Query 的配置支持和字段映射
- 完善测试覆盖：
  * 添加 TestAIQuery 单元测试，包含多种 OutputSchema 使用场景
  * 添加 TestToolAIQuery MCP 工具测试
  * 定义 GameInfo、UIElementInfo 等结构体用于测试
- 更新文档：
  * 在 docs/uixt/ai.md 中添加完整的 AIQuery 使用指南
  * 包含基本用法、OutputSchema 示例、最佳实践等
- 支持复杂的嵌套结构体和数组类型的 OutputSchema
- 与现有 AIAction、AIAssert 功能保持一致的 API 设计

uixt/driver_ext_ai.go

@@ -322,7 +322,37 @@ type SessionData struct {
 }
 
 func (dExt *XTDriver) AIQuery(text string, opts ...option.ActionOption) (string, error) {
-	return "", nil
+	if dExt.LLMService == nil {
+		return "", errors.New("LLM service is not initialized")
+	}
+
+	screenShotBase64, err := GetScreenShotBufferBase64(dExt.IDriver)
+	if err != nil {
+		return "", err
+	}
+
+	// get window size
+	size, err := dExt.IDriver.WindowSize()
+	if err != nil {
+		return "", errors.Wrap(err, "get window size for AI query failed")
+	}
+
+	// parse action options to extract OutputSchema
+	actionOptions := option.NewActionOptions(opts...)
+
+	// execute query
+	queryOpts := &ai.QueryOptions{
+		Query:        text,
+		Screenshot:   screenShotBase64,
+		Size:         size,
+		OutputSchema: actionOptions.OutputSchema,
+	}
+	result, err := dExt.LLMService.Query(context.Background(), queryOpts)
+	if err != nil {
+		return "", errors.Wrap(err, "AI query failed")
+	}
+
+	return result.Content, nil
 }
 
 func (dExt *XTDriver) AIAssert(assertion string, opts ...option.ActionOption) error {

uixt/mcp_server.go

@@ -127,6 +127,7 @@ func (s *MCPServer4XTDriver) registerTools() {
 	// AI Tools
 	s.registerTool(&ToolStartToGoal{})
 	s.registerTool(&ToolAIAction{})
+	s.registerTool(&ToolAIQuery{})
 	s.registerTool(&ToolFinished{})
 }
 

uixt/mcp_server_test.go

@@ -115,6 +115,7 @@ func TestToolInterfaces(t *testing.T) {
 		&ToolSecondaryClickBySelector{},
 		&ToolWebCloseTab{},
 		&ToolAIAction{},
+		&ToolAIQuery{},
 		&ToolFinished{},
 	}
 
@@ -1308,6 +1309,39 @@ func TestToolAIAction(t *testing.T) {
 	assert.Error(t, err)
 }
 
+// TestToolAIQuery tests the ToolAIQuery implementation
+func TestToolAIQuery(t *testing.T) {
+	tool := &ToolAIQuery{}
+
+	// Test Name
+	assert.Equal(t, option.ACTION_Query, tool.Name())
+
+	// Test Description
+	assert.NotEmpty(t, tool.Description())
+
+	// Test Options
+	options := tool.Options()
+	assert.NotNil(t, options)
+
+	// Test ConvertActionToCallToolRequest with valid params
+	action := option.MobileAction{
+		Method: option.ACTION_Query,
+		Params: "What is displayed on the screen?",
+	}
+	request, err := tool.ConvertActionToCallToolRequest(action)
+	assert.NoError(t, err)
+	assert.Equal(t, string(option.ACTION_Query), request.Params.Name)
+	assert.Equal(t, "What is displayed on the screen?", request.Params.Arguments["prompt"])
+
+	// Test ConvertActionToCallToolRequest with invalid params
+	invalidAction := option.MobileAction{
+		Method: option.ACTION_Query,
+		Params: 123, // should be string
+	}
+	_, err = tool.ConvertActionToCallToolRequest(invalidAction)
+	assert.Error(t, err)
+}
+
 // TestToolFinished tests the ToolFinished implementation
 func TestToolFinished(t *testing.T) {
 	tool := &ToolFinished{}

uixt/mcp_tools_ai.go

@@ -130,6 +130,71 @@ func (t *ToolAIAction) ConvertActionToCallToolRequest(action option.MobileAction
 	return mcp.CallToolRequest{}, fmt.Errorf("invalid AI action params: %v", action.Params)
 }
 
+// ToolAIQuery implements the ai_query tool call.
+type ToolAIQuery struct {
+	// Return data fields - these define the structure of data returned by this tool
+	Prompt string `json:"prompt" desc:"AI query prompt that was executed"`
+	Result string `json:"result" desc:"Query result content"`
+}
+
+func (t *ToolAIQuery) Name() option.ActionName {
+	return option.ACTION_Query
+}
+
+func (t *ToolAIQuery) Description() string {
+	return "Query information from screen using AI vision model with natural language prompts"
+}
+
+func (t *ToolAIQuery) Options() []mcp.ToolOption {
+	unifiedReq := &option.ActionOptions{}
+	return unifiedReq.GetMCPOptions(option.ACTION_Query)
+}
+
+func (t *ToolAIQuery) Implement() server.ToolHandlerFunc {
+	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		driverExt, err := setupXTDriver(ctx, request.Params.Arguments)
+		if err != nil {
+			return nil, fmt.Errorf("setup driver failed: %w", err)
+		}
+
+		unifiedReq, err := parseActionOptions(request.Params.Arguments)
+		if err != nil {
+			return nil, err
+		}
+
+		// Build action options from unified request
+		opts := unifiedReq.Options()
+
+		// AI query logic with options
+		result, err := driverExt.AIQuery(unifiedReq.Prompt, opts...)
+		if err != nil {
+			return NewMCPErrorResponse(fmt.Sprintf("AI query failed: %s", err.Error())), nil
+		}
+
+		message := fmt.Sprintf("Successfully queried information with prompt: %s", unifiedReq.Prompt)
+		returnData := ToolAIQuery{
+			Prompt: unifiedReq.Prompt,
+			Result: result,
+		}
+
+		return NewMCPSuccessResponse(message, &returnData), nil
+	}
+}
+
+func (t *ToolAIQuery) ConvertActionToCallToolRequest(action option.MobileAction) (mcp.CallToolRequest, error) {
+	if prompt, ok := action.Params.(string); ok {
+		arguments := map[string]any{
+			"prompt": prompt,
+		}
+
+		// Extract options to arguments
+		extractActionOptionsToArguments(action.GetOptions(), arguments)
+
+		return buildMCPCallToolRequest(t.Name(), arguments), nil
+	}
+	return mcp.CallToolRequest{}, fmt.Errorf("invalid AI query params: %v", action.Params)
+}
+
 // ToolFinished implements the finished tool call.
 type ToolFinished struct {
 	// Return data fields - these define the structure of data returned by this tool

uixt/option/action.go

@@ -184,11 +184,12 @@ type ActionOptions struct {
 	Params []float64 `json:"params,omitempty" yaml:"params,omitempty" desc:"Generic parameter array"`
 
 	// AI related
-	Prompt       string `json:"prompt,omitempty" yaml:"prompt,omitempty" desc:"AI action prompt"`
-	Content      string `json:"content,omitempty" yaml:"content,omitempty" desc:"Content for finished action"`
-	LLMService   string `json:"llm_service,omitempty" yaml:"llm_service,omitempty" desc:"LLM service type for AI actions"`
-	CVService    string `json:"cv_service,omitempty" yaml:"cv_service,omitempty" desc:"Computer vision service type for AI actions"`
-	ResetHistory bool   `json:"reset_history,omitempty" yaml:"reset_history,omitempty" desc:"Whether to reset conversation history before AI planning"`
+	Prompt       string      `json:"prompt,omitempty" yaml:"prompt,omitempty" desc:"AI action prompt"`
+	Content      string      `json:"content,omitempty" yaml:"content,omitempty" desc:"Content for finished action"`
+	LLMService   string      `json:"llm_service,omitempty" yaml:"llm_service,omitempty" desc:"LLM service type for AI actions"`
+	CVService    string      `json:"cv_service,omitempty" yaml:"cv_service,omitempty" desc:"Computer vision service type for AI actions"`
+	ResetHistory bool        `json:"reset_history,omitempty" yaml:"reset_history,omitempty" desc:"Whether to reset conversation history before AI planning"`
+	OutputSchema interface{} `json:"output_schema,omitempty" yaml:"output_schema,omitempty" desc:"Custom output schema for structured AI query response"`
 
 	// Time related
 	Seconds      float64 `json:"seconds,omitempty" yaml:"seconds,omitempty" desc:"Sleep duration in seconds"`
@@ -558,6 +559,13 @@ func WithResetHistory(resetHistory bool) ActionOption {
 	}
 }
 
+// WithOutputSchema sets the custom output schema for structured AI query response
+func WithOutputSchema(schema interface{}) ActionOption {
+	return func(o *ActionOptions) {
+		o.OutputSchema = schema
+	}
+}
+
 // HTTP API direct usage methods
 
 // ValidateForHTTPAPI validates the request for HTTP API usage
@@ -700,6 +708,9 @@ func (o *ActionOptions) validateActionSpecificFields(actionType ActionName) erro
 		ACTION_StartToGoal: func() error {
 			return o.requireFields("prompt", o.Prompt != "")
 		},
+		ACTION_Query: func() error {
+			return o.requireFields("prompt", o.Prompt != "")
+		},
 		ACTION_Finished: func() error {
 			return o.requireFields("content", o.Content != "")
 		},
@@ -774,6 +785,8 @@ func (o *ActionOptions) GetMCPOptions(actionType ActionName) []mcp.ToolOption {
 		ACTION_SleepRandom:              {"platform", "serial", "params"},
 		ACTION_AIAction:                 {"platform", "serial", "prompt", "llm_service", "cv_service"},
 		ACTION_StartToGoal:              {"platform", "serial", "prompt", "llm_service", "cv_service"},
+		ACTION_Query:                    {"platform", "serial", "prompt", "llm_service", "cv_service", "output_schema"},
+		ACTION_AIAssert:                 {"platform", "serial", "prompt", "llm_service", "cv_service"},
 		ACTION_Finished:                 {"content"},
 		ACTION_ListAvailableDevices:     {},
 		ACTION_SelectDevice:             {"platform", "serial"},
@@ -862,7 +875,15 @@ func (o *ActionOptions) generateMCPOptionsForFields(fields []string) []mcp.ToolO
 				}
 			}
 		case reflect.Map, reflect.Interface:
-			// Skip map and interface types for now
+			// Handle OutputSchema as object type
+			if name == "output_schema" {
+				if required {
+					options = append(options, mcp.WithObject(name, mcp.Required(), mcp.Description(desc)))
+				} else {
+					options = append(options, mcp.WithObject(name, mcp.Description(desc)))
+				}
+			}
+			// Skip other map and interface types for now
 			continue
 		default:
 			log.Warn().Str("field_type", fieldType.String()).Msg("Unsupported field type")