feat: add AI Querier module with custom output schema support and refactor common model calling logic

- Add new AI Querier module for structured information extraction from screenshots - Support custom output schema for structured data response - Implement automatic type conversion and data validation - Add comprehensive test suite with various data structure examples - Refactor callModelWithLogging to utils.go as shared function for planner, asserter, and querier - Eliminate code duplication across AI modules (30+ lines of repeated code) - Improve maintainability with unified logging and timing logic - Add environment variable checks in test setup to handle missing API keys gracefully Key features: - Custom output schema support with JSON Schema generation - Automatic data type conversion with reflection - Fallback mechanisms for robust parsing - Comprehensive documentation and usage examples - Backward compatibility with existing functionality
2026-09-05 23:51:25 +08:00 · 2025-06-10 20:41:35 +08:00
parent fa9a53d2ae
commit 7c45acd061
10 changed files with 1495 additions and 22 deletions
@@ -0,0 +1,515 @@
+package ai
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+
+	"github.com/cloudwego/eino-ext/components/model/openai"
+	openai2 "github.com/cloudwego/eino-ext/libs/acl/openai"
+	"github.com/cloudwego/eino/components/model"
+	"github.com/cloudwego/eino/schema"
+	"github.com/getkin/kin-openapi/openapi3gen"
+	"github.com/httprunner/httprunner/v5/code"
+	"github.com/httprunner/httprunner/v5/internal/json"
+	"github.com/httprunner/httprunner/v5/uixt/option"
+	"github.com/httprunner/httprunner/v5/uixt/types"
+	"github.com/pkg/errors"
+)
+
+// IQuerier interface defines the contract for query operations
+type IQuerier interface {
+	Query(ctx context.Context, opts *QueryOptions) (*QueryResult, error)
+}
+
+// QueryOptions represents the input options for query
+type QueryOptions struct {
+	Query        string      `json:"query"`                  // The query text to extract information
+	Screenshot   string      `json:"screenshot"`             // Base64 encoded screenshot
+	Size         types.Size  `json:"size"`                   // Screen dimensions
+	OutputSchema interface{} `json:"outputSchema,omitempty"` // Custom output schema for structured response
+}
+
+// QueryResult represents the response from an AI query
+type QueryResult struct {
+	Content string      `json:"content"`        // The extracted content/information
+	Thought string      `json:"thought"`        // The reasoning process
+	Data    interface{} `json:"data,omitempty"` // Structured data when OutputSchema is provided
+}
+
+// Querier handles query operations using different AI models
+type Querier struct {
+	modelConfig  *ModelConfig
+	model        model.ToolCallingChatModel
+	systemPrompt string
+	history      ConversationHistory
+}
+
+// NewQuerier creates a new Querier instance
+func NewQuerier(ctx context.Context, modelConfig *ModelConfig) (*Querier, error) {
+	querier := &Querier{
+		modelConfig:  modelConfig,
+		systemPrompt: defaultQueryPrompt,
+	}
+
+	if option.IS_UI_TARS(modelConfig.ModelType) {
+		querier.systemPrompt += "\n" + uiTarsQueryResponseFormat
+	} else {
+		// define default output format
+		type OutputFormat struct {
+			Content string `json:"content"`
+			Thought string `json:"thought"`
+			Error   string `json:"error,omitempty"`
+		}
+		outputFormatSchema, err := openapi3gen.NewSchemaRefForValue(&OutputFormat{}, nil)
+		if err != nil {
+			return nil, errors.Wrap(code.LLMPrepareRequestError, err.Error())
+		}
+		// set structured response format
+		modelConfig.ChatModelConfig.ResponseFormat = &openai2.ChatCompletionResponseFormat{
+			Type: openai2.ChatCompletionResponseFormatTypeJSONSchema,
+			JSONSchema: &openai2.ChatCompletionResponseFormatJSONSchema{
+				Name:        "query_result",
+				Description: "data that describes query result",
+				Schema:      outputFormatSchema.Value,
+				Strict:      false,
+			},
+		}
+	}
+
+	var err error
+	querier.model, err = openai.NewChatModel(ctx, modelConfig.ChatModelConfig)
+	if err != nil {
+		return nil, errors.Wrap(code.LLMPrepareRequestError, err.Error())
+	}
+
+	return querier, nil
+}
+
+// callModelWithLogging calls the model with automatic logging and timing
+
+// Query performs the information extraction from the screenshot
+func (q *Querier) Query(ctx context.Context, opts *QueryOptions) (*QueryResult, error) {
+	// Validate input parameters
+	if err := validateQueryInput(opts); err != nil {
+		return nil, errors.Wrap(err, "validate query parameters failed")
+	}
+
+	// Handle custom output schema if provided
+	if opts.OutputSchema != nil {
+		return q.queryWithCustomSchema(ctx, opts)
+	}
+
+	// Reset history for each new query
+	q.history = ConversationHistory{
+		{
+			Role:    schema.System,
+			Content: q.systemPrompt,
+		},
+	}
+
+	// Create user message with screenshot and query
+	userMsg := &schema.Message{
+		Role: schema.User,
+		MultiContent: []schema.ChatMessagePart{
+			{
+				Type: schema.ChatMessagePartTypeImageURL,
+				ImageURL: &schema.ChatMessageImageURL{
+					URL:    opts.Screenshot,
+					Detail: schema.ImageURLDetailAuto,
+				},
+			},
+			{
+				Type: schema.ChatMessagePartTypeText,
+				Text: fmt.Sprintf(`
+Here is the query. Please extract the requested information from the screenshot.
+=====================================
+%s
+=====================================
+  `, opts.Query),
+			},
+		},
+	}
+
+	// Append user message to history
+	q.history.Append(userMsg)
+
+	// Call model service with logging
+	message, err := callModelWithLogging(ctx, q.model, q.history,
+		q.modelConfig.ModelType, "query")
+	if err != nil {
+		return nil, errors.Wrap(code.LLMRequestServiceError, err.Error())
+	}
+
+	// Parse result
+	result, err := parseQueryResult(message.Content)
+	if err != nil {
+		return nil, errors.Wrap(code.LLMParseQueryResponseError, err.Error())
+	}
+
+	// Append assistant message to history
+	q.history.Append(&schema.Message{
+		Role:    schema.Assistant,
+		Content: message.Content,
+	})
+
+	return result, nil
+}
+
+// validateQueryInput validates the input parameters for query
+func validateQueryInput(opts *QueryOptions) error {
+	if opts.Query == "" {
+		return errors.Wrap(code.LLMPrepareRequestError, "query text is required")
+	}
+	if opts.Screenshot == "" {
+		return errors.Wrap(code.LLMPrepareRequestError, "screenshot is required")
+	}
+	return nil
+}
+
+// parseQueryResult parses the model response into QueryResult
+func parseQueryResult(content string) (*QueryResult, error) {
+	// Extract JSON content from response
+	jsonContent := extractJSONFromContent(content)
+	if jsonContent == "" {
+		// If no JSON found, treat the entire content as the result
+		// This handles cases where the model returns plain text instead of JSON
+		return &QueryResult{
+			Content: content,
+			Thought: "Direct response from model",
+		}, nil
+	}
+
+	// Parse JSON response
+	var result QueryResult
+	if err := json.Unmarshal([]byte(jsonContent), &result); err != nil {
+		// If JSON parsing fails, treat the content as plain text result
+		return &QueryResult{
+			Content: content,
+			Thought: "Failed to parse as JSON, returning raw content",
+		}, nil
+	}
+
+	return &result, nil
+}
+
+// queryWithCustomSchema performs query with custom output schema
+func (q *Querier) queryWithCustomSchema(ctx context.Context, opts *QueryOptions) (*QueryResult, error) {
+	// Create a new model config with custom schema
+	modelConfig := *q.modelConfig
+
+	if !option.IS_UI_TARS(modelConfig.ModelType) {
+		// Generate schema from the provided output schema
+		outputFormatSchema, err := openapi3gen.NewSchemaRefForValue(opts.OutputSchema, nil)
+		if err != nil {
+			return nil, errors.Wrap(code.LLMPrepareRequestError, err.Error())
+		}
+
+		// Create custom response format with the provided schema
+		modelConfig.ChatModelConfig.ResponseFormat = &openai2.ChatCompletionResponseFormat{
+			Type: openai2.ChatCompletionResponseFormatTypeJSONSchema,
+			JSONSchema: &openai2.ChatCompletionResponseFormatJSONSchema{
+				Name:        "custom_query_result",
+				Description: "custom structured data response",
+				Schema:      outputFormatSchema.Value,
+				Strict:      false,
+			},
+		}
+	}
+
+	// Create a new model instance with custom schema
+	model, err := openai.NewChatModel(ctx, modelConfig.ChatModelConfig)
+	if err != nil {
+		return nil, errors.Wrap(code.LLMPrepareRequestError, err.Error())
+	}
+
+	// Reset history for each new query
+	systemPrompt := q.systemPrompt
+	if option.IS_UI_TARS(modelConfig.ModelType) {
+		systemPrompt += "\n" + uiTarsQueryResponseFormat
+	} else {
+		// Add instruction for structured output
+		systemPrompt += "\n\nPlease respond with structured data according to the specified schema. Include both the structured data and your reasoning process."
+	}
+
+	history := ConversationHistory{
+		{
+			Role:    schema.System,
+			Content: systemPrompt,
+		},
+	}
+
+	// Create user message with screenshot and query
+	userMsg := &schema.Message{
+		Role: schema.User,
+		MultiContent: []schema.ChatMessagePart{
+			{
+				Type: schema.ChatMessagePartTypeImageURL,
+				ImageURL: &schema.ChatMessageImageURL{
+					URL:    opts.Screenshot,
+					Detail: schema.ImageURLDetailAuto,
+				},
+			},
+			{
+				Type: schema.ChatMessagePartTypeText,
+				Text: fmt.Sprintf(`
+Here is the query. Please extract the requested information from the screenshot and return it in the specified structured format.
+=====================================
+%s
+=====================================
+  `, opts.Query),
+			},
+		},
+	}
+
+	// Append user message to history
+	history.Append(userMsg)
+
+	// Call model service with logging
+	message, err := callModelWithLogging(ctx, model, history, modelConfig.ModelType, "custom schema query")
+	if err != nil {
+		return nil, errors.Wrap(code.LLMRequestServiceError, err.Error())
+	}
+
+	// Parse result with custom schema
+	result, err := parseCustomSchemaResult(message.Content, opts.OutputSchema)
+	if err != nil {
+		return nil, errors.Wrap(code.LLMParseQueryResponseError, err.Error())
+	}
+
+	// Append assistant message to history
+	q.history.Append(&schema.Message{
+		Role:    schema.Assistant,
+		Content: message.Content,
+	})
+
+	return result, nil
+}
+
+// setDefaultFieldValue sets a default value for a field in the structured data using reflection
+func setDefaultFieldValue(structValue reflect.Value, fieldName, defaultValue string) {
+	if field := structValue.FieldByName(fieldName); field.IsValid() && field.CanSet() && field.Kind() == reflect.String {
+		field.SetString(defaultValue)
+	}
+}
+
+// ensureDefaultValues ensures that Content and Thought fields have default values if empty
+func ensureDefaultValues(result *QueryResult, structuredData interface{}) {
+	const (
+		defaultContent = "Structured data extracted successfully"
+		defaultThought = "Parsed structured response according to custom schema"
+	)
+
+	// Set defaults for QueryResult
+	if result.Content == "" {
+		result.Content = defaultContent
+	}
+	if result.Thought == "" {
+		result.Thought = defaultThought
+	}
+
+	// Set defaults in structured data if it's a pointer to struct
+	if structuredData != nil {
+		if structValue := reflect.ValueOf(structuredData); structValue.Kind() == reflect.Ptr {
+			if elem := structValue.Elem(); elem.IsValid() && elem.Kind() == reflect.Struct {
+				if result.Content == defaultContent {
+					setDefaultFieldValue(elem, "Content", defaultContent)
+				}
+				if result.Thought == defaultThought {
+					setDefaultFieldValue(elem, "Thought", defaultThought)
+				}
+			}
+		}
+	}
+}
+
+// parseCustomSchemaResult parses the model response with custom schema
+func parseCustomSchemaResult(content string, outputSchema interface{}) (*QueryResult, error) {
+	// Extract JSON content from response
+	jsonContent := extractJSONFromContent(content)
+	if jsonContent == "" {
+		// If no JSON found, treat the entire content as the result
+		return &QueryResult{
+			Content: content,
+			Thought: "Direct response from model",
+		}, nil
+	}
+
+	// Create a new instance of the same type as outputSchema
+	schemaType := reflect.TypeOf(outputSchema)
+	if schemaType.Kind() == reflect.Ptr {
+		schemaType = schemaType.Elem()
+	}
+
+	// Create a new instance of the schema type
+	newInstance := reflect.New(schemaType).Interface()
+
+	// Try to unmarshal directly into the schema type
+	if err := json.Unmarshal([]byte(jsonContent), newInstance); err == nil {
+		// Successfully parsed into the expected schema type
+		result := &QueryResult{
+			Data: newInstance, // Store the typed pointer directly
+		}
+
+		// Try to extract content and thought if the schema has these fields
+		schemaValue := reflect.ValueOf(newInstance).Elem()
+		if contentField := schemaValue.FieldByName("Content"); contentField.IsValid() && contentField.Kind() == reflect.String {
+			result.Content = contentField.String()
+		}
+		if thoughtField := schemaValue.FieldByName("Thought"); thoughtField.IsValid() && thoughtField.Kind() == reflect.String {
+			result.Thought = thoughtField.String()
+		}
+
+		// If no standard fields found, try to extract from map representation
+		if result.Content == "" && result.Thought == "" {
+			var dataMap map[string]interface{}
+			if err := json.Unmarshal([]byte(jsonContent), &dataMap); err == nil {
+				if content, exists := dataMap["content"]; exists {
+					if contentStr, ok := content.(string); ok {
+						result.Content = contentStr
+					}
+				}
+				if thought, exists := dataMap["thought"]; exists {
+					if thoughtStr, ok := thought.(string); ok {
+						result.Thought = thoughtStr
+					}
+				}
+			}
+		}
+
+		// Ensure default values are set
+		ensureDefaultValues(result, newInstance)
+		return result, nil
+	}
+
+	// Fallback: try to parse as generic map and then convert
+	var structuredData interface{}
+	if err := json.Unmarshal([]byte(jsonContent), &structuredData); err == nil {
+		// Try to convert the generic data to the expected schema type
+		if convertedData, err := convertToSchemaType(structuredData, outputSchema); err == nil {
+			result := &QueryResult{
+				Data: convertedData, // Store the converted typed data
+			}
+
+			// Extract content and thought from the original map
+			if dataMap, ok := structuredData.(map[string]interface{}); ok {
+				if content, exists := dataMap["content"]; exists {
+					if contentStr, ok := content.(string); ok {
+						result.Content = contentStr
+					}
+				}
+				if thought, exists := dataMap["thought"]; exists {
+					if thoughtStr, ok := thought.(string); ok {
+						result.Thought = thoughtStr
+					}
+				}
+			}
+
+			// Ensure default values are set
+			ensureDefaultValues(result, convertedData)
+			return result, nil
+		}
+
+		// If conversion failed, fall back to storing the generic data
+		if dataMap, ok := structuredData.(map[string]interface{}); ok {
+			result := &QueryResult{
+				Data: structuredData,
+			}
+
+			// Extract content and thought if present
+			if content, exists := dataMap["content"]; exists {
+				if contentStr, ok := content.(string); ok {
+					result.Content = contentStr
+				}
+			}
+			if thought, exists := dataMap["thought"]; exists {
+				if thoughtStr, ok := thought.(string); ok {
+					result.Thought = thoughtStr
+				}
+			}
+
+			// Ensure default values are set
+			ensureDefaultValues(result, nil)
+			return result, nil
+		}
+	}
+
+	// Fallback to treating as plain text
+	return &QueryResult{
+		Content: content,
+		Thought: "Failed to parse as structured data, returning raw content",
+	}, nil
+}
+
+// convertToSchemaType converts generic data to the specified schema type
+func convertToSchemaType(data interface{}, outputSchema interface{}) (interface{}, error) {
+	// Get the type of the output schema
+	schemaType := reflect.TypeOf(outputSchema)
+	if schemaType.Kind() == reflect.Ptr {
+		schemaType = schemaType.Elem()
+	}
+
+	// Create a new instance of the schema type
+	newInstance := reflect.New(schemaType).Interface()
+
+	// Convert via JSON marshaling/unmarshaling
+	jsonData, err := json.Marshal(data)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to marshal data to JSON")
+	}
+
+	if err := json.Unmarshal(jsonData, newInstance); err != nil {
+		return nil, errors.Wrap(err, "failed to unmarshal data to target type")
+	}
+
+	return newInstance, nil
+}
+
+// ConvertQueryResultData converts QueryResult.Data to the specified type T
+// This is a helper function for type-safe conversion of the structured data
+//
+// Note: When using QueryOptions.OutputSchema, the Data field is automatically
+// converted to the correct type, so this function is typically not needed.
+// This function is mainly useful for:
+// 1. Converting data when OutputSchema was not used
+// 2. Converting to a different type than the original OutputSchema
+// 3. Handling legacy code or edge cases
+func ConvertQueryResultData[T any](result *QueryResult) (*T, error) {
+	if result.Data == nil {
+		return nil, errors.New("no structured data available")
+	}
+
+	// If Data is already of the correct type, return it directly
+	if typedData, ok := result.Data.(*T); ok {
+		return typedData, nil
+	}
+
+	// If Data is a pointer to the correct type, dereference and return
+	if reflect.TypeOf(result.Data).Kind() == reflect.Ptr {
+		if typedData, ok := result.Data.(*T); ok {
+			return typedData, nil
+		}
+		// Try to get the value that the pointer points to
+		dataValue := reflect.ValueOf(result.Data)
+		if dataValue.Kind() == reflect.Ptr && !dataValue.IsNil() {
+			elem := dataValue.Elem()
+			if elem.Type() == reflect.TypeOf((*T)(nil)).Elem() {
+				typedData := elem.Interface().(T)
+				return &typedData, nil
+			}
+		}
+	}
+
+	// Fallback: try to convert via JSON marshaling/unmarshaling
+	jsonData, err := json.Marshal(result.Data)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to marshal data to JSON")
+	}
+
+	var converted T
+	if err := json.Unmarshal(jsonData, &converted); err != nil {
+		return nil, errors.Wrap(err, "failed to unmarshal data to target type")
+	}
+
+	return &converted, nil
+}