Files
httprunner/uixt/ai/ai.go
lilong.129 fbc888655f feat: optimize ILLMService interface to support different models for each component
- Add LLMServiceConfig to support mixed model configuration
- Enable Planner, Asserter, Querier to use different optimal models
- Provide recommended configurations for various use cases
- Maintain backward compatibility with existing API
- Update documentation to reflect current state without iteration history
- Merge test files and add comprehensive configuration tests
- Resolve circular dependency by moving config to option package
2025-06-11 12:18:31 +08:00

97 lines
3.0 KiB
Go

package ai
import (
"context"
"github.com/cloudwego/eino/schema"
"github.com/httprunner/httprunner/v5/uixt/option"
)
// ILLMService 定义了 LLM 服务接口,包括规划、断言和查询功能
type ILLMService interface {
Plan(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
}
// NewLLMService creates a new LLM service with the same model for all components (backward compatibility)
func NewLLMService(modelType option.LLMServiceType) (ILLMService, error) {
config := option.NewLLMServiceConfig(modelType)
return NewLLMServiceWithOptionConfig(config)
}
// NewLLMServiceWithOptionConfig creates a new LLM service with different models for each component
func NewLLMServiceWithOptionConfig(config *option.LLMServiceConfig) (ILLMService, error) {
// Get model configs for each component
plannerModelConfig, err := GetModelConfig(config.PlannerModel)
if err != nil {
return nil, err
}
asserterModelConfig, err := GetModelConfig(config.AsserterModel)
if err != nil {
return nil, err
}
querierModelConfig, err := GetModelConfig(config.QuerierModel)
if err != nil {
return nil, err
}
// Create components with their respective model configs
planner, err := NewPlanner(context.Background(), plannerModelConfig)
if err != nil {
return nil, err
}
asserter, err := NewAsserter(context.Background(), asserterModelConfig)
if err != nil {
return nil, err
}
querier, err := NewQuerier(context.Background(), querierModelConfig)
if err != nil {
return nil, err
}
return &combinedLLMService{
planner: planner,
asserter: asserter,
querier: querier,
}, nil
}
// combinedLLMService 实现了 ILLMService 接口,组合了规划、断言和查询功能
// ⭐️支持采用不同的模型服务进行规划、断言和查询
type combinedLLMService struct {
planner IPlanner // 提供规划功能
asserter IAsserter // 提供断言功能
querier IQuerier // 提供查询功能
}
// Plan 执行规划功能
func (c *combinedLLMService) Plan(ctx context.Context, opts *PlanningOptions) (*PlanningResult, error) {
return c.planner.Plan(ctx, opts)
}
// Assert 执行断言功能
func (c *combinedLLMService) Assert(ctx context.Context, opts *AssertOptions) (*AssertionResult, error) {
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 and querier don't need tools
if planner, ok := c.planner.(*Planner); ok {
return planner.RegisterTools(tools)
}
return nil
}