Skip to content

Commit 4c7878f

Browse files
zxh326zxh326
committed
fix ai agent
Signed-off-by: Zzde <zhangxh1997@gmail.com>
1 parent 0bcd132 commit 4c7878f

File tree

9 files changed

+509
-106
lines changed

9 files changed

+509
-106
lines changed

‎pkg/ai/anthropic.go‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ func (a *Agent) processChatAnthropic(c *gin.Context, req *ChatRequest, sendEvent
3636
messages := toAnthropicMessages(req.Messages)
3737
tools := AnthropicToolDefs()
3838

39-
maxIterations := 10
39+
maxIterations := 100
4040
for i := 0; i < maxIterations; i++ {
4141
stream := a.anthropicClient.Messages.NewStreaming(ctx, anthropic.MessageNewParams{
4242
Model: anthropic.Model(a.model),
@@ -91,8 +91,9 @@ func (a *Agent) processChatAnthropic(c *gin.Context, req *ChatRequest, sendEvent
9191
sendEvent(SSEEvent{
9292
Event: "tool_result",
9393
Data: map[string]interface{}{
94-
"tool": toolName,
95-
"result": result,
94+
"tool": toolName,
95+
"result": result,
96+
"is_error": true,
9697
},
9798
})
9899
toolResults = append(toolResults, anthropic.NewToolResultBlock(tc.ID, "Tool error: "+result, true))
@@ -113,8 +114,9 @@ func (a *Agent) processChatAnthropic(c *gin.Context, req *ChatRequest, sendEvent
113114
sendEvent(SSEEvent{
114115
Event: "tool_result",
115116
Data: map[string]interface{}{
116-
"tool": toolName,
117-
"result": result,
117+
"tool": toolName,
118+
"result": result,
119+
"is_error": isError,
118120
},
119121
})
120122

‎pkg/ai/openai.go‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ func (a *Agent) processChatOpenAI(c *gin.Context, req *ChatRequest, sendEvent fu
4040

4141
tools := OpenAIToolDefs()
4242

43-
maxIterations := 10
43+
maxIterations := 100
4444
for i := 0; i < maxIterations; i++ {
4545
stream := a.openaiClient.Chat.Completions.NewStreaming(ctx, openai.ChatCompletionNewParams{
4646
Model: a.model,
@@ -99,8 +99,9 @@ func (a *Agent) processChatOpenAI(c *gin.Context, req *ChatRequest, sendEvent fu
9999
sendEvent(SSEEvent{
100100
Event: "tool_result",
101101
Data: map[string]interface{}{
102-
"tool": toolName,
103-
"result": result,
102+
"tool": toolName,
103+
"result": result,
104+
"is_error": true,
104105
},
105106
})
106107
messages = append(messages, openai.ToolMessage("Tool error: "+result, tc.ID))
@@ -121,8 +122,9 @@ func (a *Agent) processChatOpenAI(c *gin.Context, req *ChatRequest, sendEvent fu
121122
sendEvent(SSEEvent{
122123
Event: "tool_result",
123124
Data: map[string]interface{}{
124-
"tool": toolName,
125-
"result": result,
125+
"tool": toolName,
126+
"result": result,
127+
"is_error": isError,
126128
},
127129
})
128130

‎pkg/ai/tools.go‎

Lines changed: 146 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,12 @@ import (
1818
"github.com/zxh326/kite/pkg/rbac"
1919
corev1 "k8s.io/api/core/v1"
2020
apierrors "k8s.io/apimachinery/pkg/api/errors"
21+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2122
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
2223
"k8s.io/apimachinery/pkg/labels"
2324
"k8s.io/apimachinery/pkg/runtime/schema"
2425
k8stypes "k8s.io/apimachinery/pkg/types"
26+
"k8s.io/client-go/discovery"
2527
"k8s.io/klog/v2"
2628
"sigs.k8s.io/controller-runtime/pkg/client"
2729
"sigs.k8s.io/yaml"
@@ -226,7 +228,7 @@ type resourceInfo struct {
226228
ClusterScoped bool
227229
}
228230

229-
func resolveResourceInfo(kind string) resourceInfo {
231+
func resolveStaticResourceInfo(kind string) resourceInfo {
230232
switch strings.ToLower(strings.TrimSpace(kind)) {
231233
case "pod", "pods":
232234
return resourceInfo{Kind: "Pod", Resource: "pods", Version: "v1"}
@@ -264,6 +266,8 @@ func resolveResourceInfo(kind string) resourceInfo {
264266
return resourceInfo{Kind: "NetworkPolicy", Resource: "networkpolicies", Group: "networking.k8s.io", Version: "v1"}
265267
case "storageclass", "storageclasses", "sc":
266268
return resourceInfo{Kind: "StorageClass", Resource: "storageclasses", Group: "storage.k8s.io", Version: "v1", ClusterScoped: true}
269+
case "customresourcedefinition", "customresourcedefinitions", "crd", "crds":
270+
return resourceInfo{Kind: "CustomResourceDefinition", Resource: "customresourcedefinitions", Group: "apiextensions.k8s.io", Version: "v1", ClusterScoped: true}
267271
case "event", "events":
268272
return resourceInfo{Kind: "Event", Resource: "events", Version: "v1"}
269273
default:
@@ -284,6 +288,131 @@ func resolveResourceInfo(kind string) resourceInfo {
284288
}
285289
}
286290

291+
func resolveResourceInfo(ctx context.Context, cs *cluster.ClientSet, kind string) resourceInfo {
292+
if info, ok := resolveResourceInfoFromDiscovery(ctx, cs, kind, ""); ok {
293+
return info
294+
}
295+
return resolveStaticResourceInfo(kind)
296+
}
297+
298+
func resolveResourceInfoForObject(ctx context.Context, cs *cluster.ClientSet, obj *unstructured.Unstructured) resourceInfo {
299+
if info, ok := resolveResourceInfoFromDiscovery(ctx, cs, obj.GetKind(), obj.GetAPIVersion()); ok {
300+
return info
301+
}
302+
return resolveStaticResourceInfo(obj.GetKind())
303+
}
304+
305+
func resolveResourceInfoFromDiscovery(ctx context.Context, cs *cluster.ClientSet, kind, apiVersion string) (resourceInfo, bool) {
306+
input := strings.ToLower(strings.TrimSpace(kind))
307+
if input == "" || cs == nil || cs.K8sClient == nil || cs.K8sClient.ClientSet == nil {
308+
return resourceInfo{}, false
309+
}
310+
if ctx != nil {
311+
select {
312+
case <-ctx.Done():
313+
return resourceInfo{}, false
314+
default:
315+
}
316+
}
317+
discoveryClient := cs.K8sClient.ClientSet.Discovery()
318+
319+
if gv, ok := parseGroupVersion(apiVersion); ok {
320+
resourceList, err := discoveryClient.ServerResourcesForGroupVersion(gv.String())
321+
if err != nil {
322+
klog.V(2).Infof("AI tool discovery failed for %s: %v", gv.String(), err)
323+
} else if info, found := findResourceInfoInList(input, gv, resourceList.APIResources); found {
324+
return info, true
325+
}
326+
}
327+
328+
resourceLists, err := discoveryClient.ServerPreferredResources()
329+
if err != nil && !discovery.IsGroupDiscoveryFailedError(err) {
330+
klog.V(2).Infof("AI tool preferred discovery failed: %v", err)
331+
return resourceInfo{}, false
332+
}
333+
334+
for _, resourceList := range resourceLists {
335+
if resourceList == nil {
336+
continue
337+
}
338+
gv, err := schema.ParseGroupVersion(resourceList.GroupVersion)
339+
if err != nil {
340+
continue
341+
}
342+
if info, found := findResourceInfoInList(input, gv, resourceList.APIResources); found {
343+
return info, true
344+
}
345+
}
346+
347+
return resourceInfo{}, false
348+
}
349+
350+
func parseGroupVersion(apiVersion string) (schema.GroupVersion, bool) {
351+
apiVersion = strings.TrimSpace(apiVersion)
352+
if apiVersion == "" {
353+
return schema.GroupVersion{}, false
354+
}
355+
gv, err := schema.ParseGroupVersion(apiVersion)
356+
if err != nil {
357+
return schema.GroupVersion{}, false
358+
}
359+
return gv, true
360+
}
361+
362+
func findResourceInfoInList(input string, gv schema.GroupVersion, apiResources []metav1.APIResource) (resourceInfo, bool) {
363+
group := strings.ToLower(gv.Group)
364+
for _, apiResource := range apiResources {
365+
if strings.Contains(apiResource.Name, "/") {
366+
continue
367+
}
368+
if !resourceMatchesInput(input, group, apiResource) {
369+
continue
370+
}
371+
return resourceInfo{
372+
Kind: apiResource.Kind,
373+
Resource: apiResource.Name,
374+
Group: gv.Group,
375+
Version: gv.Version,
376+
ClusterScoped: !apiResource.Namespaced,
377+
}, true
378+
}
379+
return resourceInfo{}, false
380+
}
381+
382+
func resourceMatchesInput(input, group string, apiResource metav1.APIResource) bool {
383+
candidates := make([]string, 0, 3+len(apiResource.ShortNames))
384+
if kind := strings.ToLower(strings.TrimSpace(apiResource.Kind)); kind != "" {
385+
candidates = append(candidates, kind)
386+
}
387+
if name := strings.ToLower(strings.TrimSpace(apiResource.Name)); name != "" {
388+
candidates = append(candidates, name)
389+
}
390+
if singular := strings.ToLower(strings.TrimSpace(apiResource.SingularName)); singular != "" {
391+
candidates = append(candidates, singular)
392+
}
393+
for _, shortName := range apiResource.ShortNames {
394+
if shortName = strings.ToLower(strings.TrimSpace(shortName)); shortName != "" {
395+
candidates = append(candidates, shortName)
396+
}
397+
}
398+
399+
for _, candidate := range candidates {
400+
if input == candidate {
401+
return true
402+
}
403+
if !strings.HasSuffix(candidate, "s") && input == candidate+"s" {
404+
return true
405+
}
406+
if group != "" && input == candidate+"."+group {
407+
return true
408+
}
409+
if group != "" && !strings.HasSuffix(candidate, "s") && input == candidate+"s."+group {
410+
return true
411+
}
412+
}
413+
return false
414+
}
415+
287416
func (r resourceInfo) GVK() schema.GroupVersionKind {
288417
return schema.GroupVersionKind{Group: r.Group, Version: r.Version, Kind: r.Kind}
289418
}
@@ -299,8 +428,7 @@ func normalizeNamespace(r resourceInfo, namespace string) string {
299428
return namespace
300429
}
301430

302-
func buildObjectForKind(kind string) *unstructured.Unstructured {
303-
resource := resolveResourceInfo(kind)
431+
func buildObjectForResource(resource resourceInfo) *unstructured.Unstructured {
304432
obj := &unstructured.Unstructured{}
305433
obj.SetGroupVersionKind(resource.GVK())
306434
return obj
@@ -356,15 +484,15 @@ func permissionNamespace(resource resourceInfo, namespace string) string {
356484
return namespace
357485
}
358486

359-
func requiredToolPermissions(toolName string, args map[string]interface{}) ([]toolPermission, error) {
487+
func requiredToolPermissions(ctx context.Context, cs *cluster.ClientSet, toolName string, args map[string]interface{}) ([]toolPermission, error) {
360488
switch toolName {
361489
case "get_resource":
362490
kind, err := getRequiredString(args, "kind")
363491
if err != nil {
364492
return nil, err
365493
}
366494
namespace, _ := args["namespace"].(string)
367-
resource := resolveResourceInfo(kind)
495+
resource := resolveResourceInfo(ctx, cs, kind)
368496
return []toolPermission{{
369497
Resource: resource.Resource,
370498
Verb: string(common.VerbGet),
@@ -376,7 +504,7 @@ func requiredToolPermissions(toolName string, args map[string]interface{}) ([]to
376504
return nil, err
377505
}
378506
namespace, _ := args["namespace"].(string)
379-
resource := resolveResourceInfo(kind)
507+
resource := resolveResourceInfo(ctx, cs, kind)
380508
return []toolPermission{{
381509
Resource: resource.Resource,
382510
Verb: string(common.VerbGet),
@@ -407,7 +535,7 @@ func requiredToolPermissions(toolName string, args map[string]interface{}) ([]to
407535
if err != nil {
408536
return nil, err
409537
}
410-
resource := resolveResourceInfo(obj.GetKind())
538+
resource := resolveResourceInfoForObject(ctx, cs, obj)
411539
return []toolPermission{{
412540
Resource: resource.Resource,
413541
Verb: string(common.VerbCreate),
@@ -418,7 +546,7 @@ func requiredToolPermissions(toolName string, args map[string]interface{}) ([]to
418546
if err != nil {
419547
return nil, err
420548
}
421-
resource := resolveResourceInfo(obj.GetKind())
549+
resource := resolveResourceInfoForObject(ctx, cs, obj)
422550
return []toolPermission{{
423551
Resource: resource.Resource,
424552
Verb: string(common.VerbUpdate),
@@ -433,7 +561,7 @@ func requiredToolPermissions(toolName string, args map[string]interface{}) ([]to
433561
return nil, err
434562
}
435563
namespace, _ := args["namespace"].(string)
436-
resource := resolveResourceInfo(kind)
564+
resource := resolveResourceInfo(ctx, cs, kind)
437565
return []toolPermission{{
438566
Resource: resource.Resource,
439567
Verb: string(common.VerbUpdate),
@@ -448,7 +576,7 @@ func requiredToolPermissions(toolName string, args map[string]interface{}) ([]to
448576
return nil, err
449577
}
450578
namespace, _ := args["namespace"].(string)
451-
resource := resolveResourceInfo(kind)
579+
resource := resolveResourceInfo(ctx, cs, kind)
452580
return []toolPermission{{
453581
Resource: resource.Resource,
454582
Verb: string(common.VerbDelete),
@@ -480,7 +608,7 @@ func AuthorizeTool(c *gin.Context, cs *cluster.ClientSet, toolName string, args
480608
return "Error: authenticated user not found in context", true
481609
}
482610

483-
permissions, err := requiredToolPermissions(toolName, args)
611+
permissions, err := requiredToolPermissions(c.Request.Context(), cs, toolName, args)
484612
if err != nil {
485613
return "Error: " + err.Error(), true
486614
}
@@ -533,8 +661,8 @@ func executeGetResource(ctx context.Context, cs *cluster.ClientSet, args map[str
533661
}
534662
namespace, _ := args["namespace"].(string)
535663

536-
resource := resolveResourceInfo(kind)
537-
obj := buildObjectForKind(kind)
664+
resource := resolveResourceInfo(ctx, cs, kind)
665+
obj := buildObjectForResource(resource)
538666
key := k8stypes.NamespacedName{
539667
Name: name,
540668
Namespace: normalizeNamespace(resource, namespace),
@@ -596,7 +724,7 @@ func executeListResources(ctx context.Context, cs *cluster.ClientSet, args map[s
596724
namespace, _ := args["namespace"].(string)
597725
labelSelector, _ := args["label_selector"].(string)
598726

599-
resource := resolveResourceInfo(kind)
727+
resource := resolveResourceInfo(ctx, cs, kind)
600728
namespace = normalizeNamespace(resource, namespace)
601729
list := &unstructured.UnstructuredList{}
602730
list.SetGroupVersionKind(resource.ListGVK())
@@ -1072,8 +1200,8 @@ func executePatchResource(ctx context.Context, cs *cluster.ClientSet, args map[s
10721200
return "Error: patch must be valid JSON", true
10731201
}
10741202

1075-
resource := resolveResourceInfo(kind)
1076-
obj := buildObjectForKind(kind)
1203+
resource := resolveResourceInfo(ctx, cs, kind)
1204+
obj := buildObjectForResource(resource)
10771205

10781206
key := k8stypes.NamespacedName{
10791207
Name: name,
@@ -1104,8 +1232,8 @@ func executeDeleteResource(ctx context.Context, cs *cluster.ClientSet, args map[
11041232
}
11051233
namespace, _ := args["namespace"].(string)
11061234

1107-
resource := resolveResourceInfo(kind)
1108-
obj := buildObjectForKind(kind)
1235+
resource := resolveResourceInfo(ctx, cs, kind)
1236+
obj := buildObjectForResource(resource)
11091237

11101238
key := k8stypes.NamespacedName{
11111239
Name: name,

0 commit comments

Comments
 (0)