From 2b32e17091a82919537a9ed03365d73353f8abdf Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Tue, 1 Jul 2025 22:39:47 +0800
Subject: [PATCH] refactor: init CV service

---
 uixt/driver_ext_screenshot.go |  4 ++++
 uixt/sdk.go                   | 35 ++++++++++++++++++++++-------------
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/uixt/driver_ext_screenshot.go b/uixt/driver_ext_screenshot.go
index 002bec2a..5ac07764 100644
--- a/uixt/driver_ext_screenshot.go
+++ b/uixt/driver_ext_screenshot.go
@@ -114,6 +114,10 @@ func (dExt *XTDriver) createScreenshotWithSession(opts ...option.ActionOption) (
 	logger := log.Debug().Str("imagePath", imagePath)
 	// perform CV processing if any CV-related option is enabled
 	if needsCVProcessing(screenshotOptions) {
+		if err = dExt.initCVService(); err != nil {
+			return nil, err
+		}
+
 		imageResult, err := dExt.CVService.ReadFromBuffer(compressBufSource, opts...)
 		if err != nil {
 			log.Error().Err(err).Msg("ReadFromBuffer from ImageService failed")
diff --git a/uixt/sdk.go b/uixt/sdk.go
index ad23c3c5..d315bbca 100644
--- a/uixt/sdk.go
+++ b/uixt/sdk.go
@@ -15,29 +15,18 @@ import (
 )
 
 func NewXTDriver(driver IDriver, opts ...option.AIServiceOption) (*XTDriver, error) {
+	services := option.NewAIServiceOptions(opts...)
 	driverExt := &XTDriver{
 		IDriver: driver,
 		client: &MCPClient4XTDriver{
 			Server: NewMCPServer(),
 		},
+		services:         services,
 		loadedMCPClients: make(map[string]client.MCPClient),
 	}
 
-	services := option.NewAIServiceOptions(opts...)
-
 	var err error
 
-	// default to vedem CV service
-	if services.CVService == "" {
-		log.Warn().Msg("no CV service config provided, not using CV service")
-	} else {
-		driverExt.CVService, err = ai.NewCVService(services.CVService)
-		if err != nil {
-			log.Error().Err(err).Msg("init vedem image service failed")
-			return nil, err
-		}
-	}
-
 	// Handle LLM service initialization
 	if services.LLMConfig != nil {
 		// Use advanced LLM configuration if provided
@@ -73,10 +62,30 @@ type XTDriver struct {
 	CVService  ai.ICVService  // OCR/CV
 	LLMService ai.ILLMService // LLM
 
+	services         *option.AIServiceOptions    // AI services options
 	client           *MCPClient4XTDriver         // MCP Client for built-in uixt server
 	loadedMCPClients map[string]client.MCPClient // External MCP clients
 }
 
+func (dExt *XTDriver) initCVService() error {
+	if dExt.CVService != nil {
+		return nil
+	}
+	cvServiceType := dExt.services.CVService
+	if cvServiceType == "" {
+		log.Warn().Msg("no CV service config provided, use default vedem")
+		cvServiceType = option.CVServiceTypeVEDEM
+	}
+	cvService, err := ai.NewCVService(cvServiceType)
+	if err != nil {
+		log.Error().Err(err).Str("type", string(cvServiceType)).
+			Msg("init cv service failed")
+		return errors.Wrap(err, "init cv service failed")
+	}
+	dExt.CVService = cvService
+	return nil
+}
+
 // MCPClient4XTDriver is a mock MCP client that only implements the methods used by the host
 type MCPClient4XTDriver struct {
 	client.MCPClient