wlw/ai-agent
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: refactor prepare service to use toolCatalog for tool resolution and add unit tests for tool catalog functionality

Browse Source
This commit is contained in:
mlogclub
2026-04-13 19:20:14 +08:00
parent 2effb35af3
commit 1ebbdadc1b
4 changed files with 198 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/ai/runtime/app/prepare_service.go
+13 -99
View File
@@ -5,18 +5,16 @@ import (
"encoding/json"
"strings"
"cs-agent/internal/ai/runtime/registry"
"cs-agent/internal/ai/skills"
"cs-agent/internal/models"
"cs-agent/internal/pkg/toolx"
)
func newPrepareService(registry *registry.Registry) *prepareService {
return &prepareService{registry: registry}
func newPrepareService(catalog *toolCatalog) *prepareService {
return &prepareService{catalog: catalog}
}
type prepareService struct {
registry *registry.Registry
catalog *toolCatalog
}
func (s *prepareService) selectSkill(ctx context.Context, req Request) (*models.SkillDefinition, string, string, error) {
@@ -44,37 +42,30 @@ func (s *prepareService) selectSkill(ctx context.Context, req Request) (*models.
}
func (s *prepareService) prepareToolsForRun(req *Request) error {
if req == nil || req.ToolSet != nil || s.registry == nil {
if req == nil || req.ToolSet != nil || s.catalog == nil {
return nil
}
toolSet, err := s.registry.Resolve(registry.Context{
Conversation: req.Conversation,
AIAgent: req.AIAgent,
AIConfig: req.AIConfig,
UserMessage: req.UserMessage,
AllowedToolCodes: resolveAllowedToolCodes(req.AIAgent, req.SelectedSkill),
})
toolSet, err := s.catalog.resolveForRun(req)
if err != nil {
return err
}
req.ToolSet = toolSet
if toolSet != nil {
req.ToolSet = toolSet
}
return nil
}
func (s *prepareService) prepareToolsForResume(req *ResumeRequest) error {
if req == nil || req.ToolSet != nil || s.registry == nil {
if req == nil || req.ToolSet != nil || s.catalog == nil {
return nil
}
toolSet, err := s.registry.Resolve(registry.Context{
Conversation: req.Conversation,
AIAgent: req.AIAgent,
AIConfig: req.AIConfig,
AllowedToolCodes: parseAgentAllowedToolCodes(req.AIAgent),
})
toolSet, err := s.catalog.resolveForResume(req)
if err != nil {
return err
}
req.ToolSet = toolSet
if toolSet != nil {
req.ToolSet = toolSet
}
return nil
}
@@ -96,80 +87,3 @@ func cloneSkillDefinition(item *models.SkillDefinition) *models.SkillDefinition
clone := *item
return &clone
}
func parseSkillAllowedToolCodes(skill *models.SkillDefinition) []string {
if skill == nil {
return nil
}
raw := strings.TrimSpace(skill.ToolWhitelist)
if raw == "" {
return nil
}
var items []string
if err := json.Unmarshal([]byte(raw), &items); err != nil {
return nil
}
ret := make([]string, 0, len(items))
for _, item := range items {
item = strings.TrimSpace(item)
item = toolx.NormalizeToolCodeAlias(item)
if item == "" {
continue
}
ret = append(ret, item)
}
return ret
}
func parseAgentAllowedToolCodes(aiAgent *models.AIAgent) []string {
if aiAgent == nil || strings.TrimSpace(aiAgent.AllowedMCPTools) == "" {
return nil
}
items, err := toolx.ParseAgentMCPToolsJSON(aiAgent.AllowedMCPTools)
if err != nil {
return nil
}
ret := make([]string, 0, len(items))
for _, item := range items {
toolCode := strings.TrimSpace(item.ToolCode)
toolCode = toolx.NormalizeToolCodeAlias(toolCode)
if toolCode == "" {
continue
}
ret = append(ret, toolCode)
}
return ret
}
func resolveAllowedToolCodes(aiAgent *models.AIAgent, skill *models.SkillDefinition) []string {
agentAllowed := parseAgentAllowedToolCodes(aiAgent)
skillAllowed := parseSkillAllowedToolCodes(skill)
switch {
case len(agentAllowed) == 0:
return skillAllowed
case len(skillAllowed) == 0:
return agentAllowed
default:
skillSet := make(map[string]struct{}, len(skillAllowed))
for _, item := range skillAllowed {
item = strings.TrimSpace(item)
item = toolx.NormalizeToolCodeAlias(item)
if item == "" {
continue
}
skillSet[item] = struct{}{}
}
ret := make([]string, 0, len(agentAllowed))
for _, item := range agentAllowed {
item = strings.TrimSpace(item)
item = toolx.NormalizeToolCodeAlias(item)
if item == "" {
continue
}
if _, ok := skillSet[item]; ok {
ret = append(ret, item)
}
}
return ret
}
}
internal/ai/runtime/app/service.go
+6 -21
View File
@@ -4,38 +4,24 @@ import (
"context"
"cs-agent/internal/ai/runtime/internal/executor"
"cs-agent/internal/ai/runtime/registry"
"cs-agent/internal/ai/runtime/tools"
)
type Service struct {
runtime *executor.Service
registry *registry.Registry
prepare *prepareService
runtime *executor.Service
catalog *toolCatalog
prepare *prepareService
}
func NewService() *Service {
catalog := newToolCatalog()
return &Service{
runtime: executor.NewService(),
registry: registry.NewRegistry(
tools.NewTriageServiceRequestTool(),
tools.NewAnalyzeConversationTool(),
tools.NewPrepareTicketDraftTool(),
tools.NewCreateTicketGraphTool(),
tools.NewHandoffGraphTool(),
),
}
}
// TODO 这个方法真的要这样吗？ 不能直接在NewService中直接初始化吗？
func (s *Service) initPrepareService() {
if s.prepare == nil {
s.prepare = newPrepareService(s.registry)
catalog: catalog,
prepare: newPrepareService(catalog),
}
}
func (s *Service) Run(ctx context.Context, req Request) (*Summary, error) {
s.initPrepareService()
selectedSkill, skillReason, skillTrace, skillErr := s.prepare.selectSkill(ctx, req)
req.SelectedSkill = selectedSkill
req.SkillRouteReason = skillReason
@@ -72,7 +58,6 @@ func (s *Service) Run(ctx context.Context, req Request) (*Summary, error) {
}
func (s *Service) Resume(ctx context.Context, req ResumeRequest) (*Summary, error) {
s.initPrepareService()
if err := s.prepare.prepareToolsForResume(&req); err != nil {
return nil, err
}
internal/ai/runtime/app/tool_catalog.go
+125
View File
@@ -0,0 +1,125 @@
package app
import (
"encoding/json"
"strings"
"cs-agent/internal/ai/runtime/registry"
"cs-agent/internal/ai/runtime/tools"
"cs-agent/internal/models"
"cs-agent/internal/pkg/toolx"
)
type toolCatalog struct {
registry *registry.Registry
}
func newToolCatalog() *toolCatalog {
return &toolCatalog{
registry: registry.NewRegistry(
tools.NewTriageServiceRequestTool(),
tools.NewAnalyzeConversationTool(),
tools.NewPrepareTicketDraftTool(),
tools.NewCreateTicketGraphTool(),
tools.NewHandoffGraphTool(),
),
}
}
func (c *toolCatalog) resolveForRun(req *Request) (*registry.ToolSet, error) {
if req == nil || req.ToolSet != nil || c == nil || c.registry == nil {
return nil, nil
}
return c.registry.Resolve(registry.Context{
Conversation: req.Conversation,
AIAgent: req.AIAgent,
AIConfig: req.AIConfig,
UserMessage: req.UserMessage,
AllowedToolCodes: c.resolveAllowedToolCodes(req.AIAgent, req.SelectedSkill),
})
}
func (c *toolCatalog) resolveForResume(req *ResumeRequest) (*registry.ToolSet, error) {
if req == nil || req.ToolSet != nil || c == nil || c.registry == nil {
return nil, nil
}
return c.registry.Resolve(registry.Context{
Conversation: req.Conversation,
AIAgent: req.AIAgent,
AIConfig: req.AIConfig,
AllowedToolCodes: c.parseAgentAllowedToolCodes(req.AIAgent),
})
}
func (c *toolCatalog) parseSkillAllowedToolCodes(skill *models.SkillDefinition) []string {
if skill == nil {
return nil
}
raw := strings.TrimSpace(skill.ToolWhitelist)
if raw == "" {
return nil
}
var items []string
if err := json.Unmarshal([]byte(raw), &items); err != nil {
return nil
}
return normalizeAllowedToolCodes(items)
}
func (c *toolCatalog) parseAgentAllowedToolCodes(aiAgent *models.AIAgent) []string {
if aiAgent == nil || strings.TrimSpace(aiAgent.AllowedMCPTools) == "" {
return nil
}
items, err := toolx.ParseAgentMCPToolsJSON(aiAgent.AllowedMCPTools)
if err != nil {
return nil
}
ret := make([]string, 0, len(items))
for _, item := range items {
ret = append(ret, item.ToolCode)
}
return normalizeAllowedToolCodes(ret)
}
func (c *toolCatalog) resolveAllowedToolCodes(aiAgent *models.AIAgent, skill *models.SkillDefinition) []string {
agentAllowed := c.parseAgentAllowedToolCodes(aiAgent)
skillAllowed := c.parseSkillAllowedToolCodes(skill)
switch {
case len(agentAllowed) == 0:
return skillAllowed
case len(skillAllowed) == 0:
return agentAllowed
default:
skillSet := make(map[string]struct{}, len(skillAllowed))
for _, item := range skillAllowed {
skillSet[item] = struct{}{}
}
ret := make([]string, 0, len(agentAllowed))
for _, item := range agentAllowed {
if _, ok := skillSet[item]; ok {
ret = append(ret, item)
}
}
return ret
}
}
func normalizeAllowedToolCodes(items []string) []string {
if len(items) == 0 {
return nil
}
ret := make([]string, 0, len(items))
seen := make(map[string]struct{}, len(items))
for _, item := range items {
item = toolx.NormalizeToolCodeAlias(strings.TrimSpace(item))
if item == "" {
continue
}
if _, ok := seen[item]; ok {
continue
}
seen[item] = struct{}{}
ret = append(ret, item)
}
return ret
}
internal/ai/runtime/app/tool_catalog_test.go
+54
View File
@@ -0,0 +1,54 @@
package app
import (
"testing"
"cs-agent/internal/models"
)
func TestNormalizeAllowedToolCodes(t *testing.T) {
ret := normalizeAllowedToolCodes([]string{
" ",
"graph/create_ticket_with_confirmation",
"builtin/create_ticket_with_confirmation",
"graph/handoff_to_human",
"graph/handoff_to_human",
})
if len(ret) != 2 {
t.Fatalf("expected 2 tool codes, got %d: %#v", len(ret), ret)
}
if ret[0] != "graph/create_ticket_with_confirmation" {
t.Fatalf("unexpected first tool code: %s", ret[0])
}
if ret[1] != "graph/handoff_to_human" {
t.Fatalf("unexpected second tool code: %s", ret[1])
}
}
func TestToolCatalogResolveAllowedToolCodes(t *testing.T) {
catalog := newToolCatalog()
agent := &models.AIAgent{
AllowedMCPTools: `[{"toolCode":"graph/create_ticket_with_confirmation"},{"toolCode":"graph/handoff_to_human"}]`,
}
skill := &models.SkillDefinition{
ToolWhitelist: `["builtin/create_ticket_with_confirmation","graph/prepare_ticket_draft"]`,
}
ret := catalog.resolveAllowedToolCodes(agent, skill)
if len(ret) != 1 {
t.Fatalf("expected 1 tool code, got %d: %#v", len(ret), ret)
}
if ret[0] != "graph/create_ticket_with_confirmation" {
t.Fatalf("unexpected tool code: %s", ret[0])
}
}
func TestToolCatalogResolveAllowedToolCodesFallsBackWhenSkillEmpty(t *testing.T) {
catalog := newToolCatalog()
agent := &models.AIAgent{
AllowedMCPTools: `[{"toolCode":"graph/create_ticket_with_confirmation"},{"toolCode":"graph/handoff_to_human"}]`,
}
ret := catalog.resolveAllowedToolCodes(agent, nil)
if len(ret) != 2 {
t.Fatalf("expected 2 tool codes, got %d: %#v", len(ret), ret)
}
}