From c513e56d30dcf1cb9188cf74533d34256cca71a9 Mon Sep 17 00:00:00 2001 From: "lilong.129" Date: Tue, 10 Jun 2025 20:45:49 +0800 Subject: [PATCH] feat: add Query method to ILLMService interface - Add Query method to ILLMService interface for unified AI service access - Update combinedLLMService to include querier functionality - Add comprehensive tests for ILLMService Query method - Support both basic query and custom schema query through unified interface - Add environment variable checks for test reliability This allows users to access all AI capabilities (planning, assertion, and query) through a single ILLMService interface, providing better API consistency and ease of use. --- internal/version/VERSION | 2 +- uixt/ai/ai.go | 20 ++++-- uixt/ai/ai_test.go | 142 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 159 insertions(+), 5 deletions(-) create mode 100644 uixt/ai/ai_test.go diff --git a/internal/version/VERSION b/internal/version/VERSION index cf1fe76d..6e73e2d6 100644 --- a/internal/version/VERSION +++ b/internal/version/VERSION @@ -1 +1 @@ -v5.0.0-beta-2506102041 +v5.0.0-beta-2506102045 diff --git a/uixt/ai/ai.go b/uixt/ai/ai.go index 3a433aad..3c0e230a 100644 --- a/uixt/ai/ai.go +++ b/uixt/ai/ai.go @@ -7,10 +7,11 @@ import ( "github.com/httprunner/httprunner/v5/uixt/option" ) -// ILLMService 定义了 LLM 服务接口，包括规划和断言功能 +// ILLMService 定义了 LLM 服务接口，包括规划、断言和查询功能 type ILLMService interface { Call(ctx context.Context, opts *PlanningOptions) (*PlanningResult, error) Assert(ctx context.Context, opts *AssertOptions) (*AssertionResult, error) + Query(ctx context.Context, opts *QueryOptions) (*QueryResult, error) // RegisterTools registers tools for function calling RegisterTools(tools []*schema.ToolInfo) error } @@ -29,18 +30,24 @@ func NewLLMService(modelType option.LLMServiceType) (ILLMService, error) { if err != nil { return nil, err } + querier, err := NewQuerier(context.Background(), modelConfig) + if err != nil { + return nil, err + } return &combinedLLMService{ planner: planner, asserter: asserter, + querier: querier, }, nil } -// combinedLLMService 实现了 ILLMService 接口，组合了规划和断言功能 -// ⭐️支持采用不同的模型服务进行规划和断言 +// combinedLLMService 实现了 ILLMService 接口，组合了规划、断言和查询功能 +// ⭐️支持采用不同的模型服务进行规划、断言和查询 type combinedLLMService struct { planner IPlanner // 提供规划功能 asserter IAsserter // 提供断言功能 + querier IQuerier // 提供查询功能 } // Call 执行规划功能 @@ -53,9 +60,14 @@ func (c *combinedLLMService) Assert(ctx context.Context, opts *AssertOptions) (* return c.asserter.Assert(ctx, opts) } +// Query 执行查询功能 +func (c *combinedLLMService) Query(ctx context.Context, opts *QueryOptions) (*QueryResult, error) { + return c.querier.Query(ctx, opts) +} + // RegisterTools registers tools for function calling func (c *combinedLLMService) RegisterTools(tools []*schema.ToolInfo) error { - // Only register tools to planner since asserter doesn't need tools + // Only register tools to planner since asserter and querier don't need tools if planner, ok := c.planner.(*Planner); ok { return planner.RegisterTools(tools) } diff --git a/uixt/ai/ai_test.go b/uixt/ai/ai_test.go new file mode 100644 index 00000000..2d49f2c9 --- /dev/null +++ b/uixt/ai/ai_test.go @@ -0,0 +1,142 @@ +package ai + +import ( + "context" + "os" + "testing" + + "github.com/httprunner/httprunner/v5/internal/builtin" + "github.com/httprunner/httprunner/v5/uixt/option" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// hasRequiredEnvVars checks if the required environment variables are set for testing +func hasRequiredEnvVars() bool { + // Check for OpenAI environment variables + if os.Getenv("OPENAI_BASE_URL") != "" && os.Getenv("OPENAI_API_KEY") != "" { + return true + } + // Check for GPT-4O specific environment variables + if os.Getenv("OPENAI_GPT_4O_BASE_URL") != "" && os.Getenv("OPENAI_GPT_4O_API_KEY") != "" { + return true + } + return false +} + +func TestILLMServiceQuery(t *testing.T) { + // Skip test if required environment variables are not set + if !hasRequiredEnvVars() { + t.Skip("Skipping test: required environment variables not set") + } + + // Create LLM service + service, err := NewLLMService(option.OPENAI_GPT_4O) + require.NoError(t, err) + require.NotNil(t, service) + + // Load test image + screenshot, size, err := builtin.LoadImage("testdata/llk_1.png") + require.NoError(t, err) + + // Test basic query functionality + t.Run("BasicQuery", func(t *testing.T) { + opts := &QueryOptions{ + Query: "请描述这张图片中的内容", + Screenshot: screenshot, + Size: size, + } + + result, err := service.Query(context.Background(), opts) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotEmpty(t, result.Content) + assert.NotEmpty(t, result.Thought) + assert.Nil(t, result.Data) // Should be nil for standard query + + t.Logf("Query result: %s", result.Content) + }) + + // Test custom schema query + t.Run("CustomSchemaQuery", func(t *testing.T) { + type GameInfo struct { + Content string `json:"content"` + Thought string `json:"thought"` + Rows int `json:"rows"` + Cols int `json:"cols"` + Icons []string `json:"icons"` + } + + opts := &QueryOptions{ + Query: "请分析这个连连看游戏界面，告诉我有多少行多少列，有哪些不同类型的图案", + Screenshot: screenshot, + Size: size, + OutputSchema: GameInfo{}, + } + + result, err := service.Query(context.Background(), opts) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotEmpty(t, result.Content) + assert.NotEmpty(t, result.Thought) + assert.NotNil(t, result.Data) + + // Verify type conversion + if gameInfo, ok := result.Data.(*GameInfo); ok { + assert.NotEmpty(t, gameInfo.Content) + assert.NotEmpty(t, gameInfo.Thought) + assert.Greater(t, gameInfo.Rows, 0) + assert.Greater(t, gameInfo.Cols, 0) + assert.NotEmpty(t, gameInfo.Icons) + t.Logf("Game info: %+v", gameInfo) + } else { + t.Errorf("Expected *GameInfo, got %T", result.Data) + } + }) +} + +func TestILLMServiceIntegration(t *testing.T) { + // Skip test if required environment variables are not set + if !hasRequiredEnvVars() { + t.Skip("Skipping test: required environment variables not set") + } + + // Create LLM service + service, err := NewLLMService(option.OPENAI_GPT_4O) + require.NoError(t, err) + require.NotNil(t, service) + + // Load test image + screenshot, size, err := builtin.LoadImage("testdata/llk_1.png") + require.NoError(t, err) + + ctx := context.Background() + + // Test that all three methods work + t.Run("AllMethods", func(t *testing.T) { + // Test Query + queryOpts := &QueryOptions{ + Query: "请分析这张图片", + Screenshot: screenshot, + Size: size, + } + queryResult, err := service.Query(ctx, queryOpts) + assert.NoError(t, err) + assert.NotNil(t, queryResult) + t.Logf("Query result: %s", queryResult.Content) + + // Test Assert + assertOpts := &AssertOptions{ + Assertion: "这是一个连连看游戏界面", + Screenshot: screenshot, + Size: size, + } + assertResult, err := service.Assert(ctx, assertOpts) + assert.NoError(t, err) + assert.NotNil(t, assertResult) + t.Logf("Assert result: pass=%v, thought=%s", assertResult.Pass, assertResult.Thought) + + // Note: Planning test would require proper user instruction and message setup + // which is more complex, so we skip it in this integration test + }) +}