fix: agent schema and tests

ferruhcihan · ferruhcihan · commit ed8120cb77a5 · 2025-10-08T16:07:29.000+02:00
diff --git a/src/ai/AkamaiAgentCR.test.ts b/src/ai/AkamaiAgentCR.test.ts
@@ -30,7 +30,12 @@ describe('AkamaiAgentCR', () => {
     spec: {
       foundationModel: 'gpt-4',
       agentInstructions: 'You are a helpful assistant',
-      knowledgeBase: 'test-kb',
+      tools: [
+        {
+          type: 'knowledgeBase',
+          name: 'test-kb',
+        },
+      ],
     },
   }
 
@@ -48,11 +53,14 @@ describe('AkamaiAgentCR', () => {
       expect(agentCR.metadata.namespace).toBe('team-team-123')
       expect(agentCR.metadata.labels?.['apl.io/teamId']).toBe('team-123')
       expect(agentCR.spec.foundationModel).toBe('gpt-4')
-      expect(agentCR.spec.systemPrompt).toBe('You are a helpful assistant')
+      expect(agentCR.spec.agentInstructions).toBe('You are a helpful assistant')
       expect(agentCR.spec.tools).toEqual([
         {
           type: 'knowledgeBase',
           name: 'test-kb',
+          description:
+            'Search the test-kb knowledge base for relevant information. Use this when you need factual information, documentation, or specific details stored in the knowledge base.',
+          endpoint: undefined,
         },
       ])
     })
@@ -67,19 +75,46 @@ describe('AkamaiAgentCR', () => {
       expect(agentCR.metadata.labels?.['custom-label']).toBeUndefined()
     })
 
-    test('should handle request without knowledgeBase', () => {
-      const requestWithoutKB = {
+    test('should handle request without tools', () => {
+      const requestWithoutTools = {
         ...mockAgentRequest,
         spec: {
           ...mockAgentRequest.spec,
-          knowledgeBase: undefined,
+          tools: undefined,
         },
       }
 
-      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutKB)
+      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutTools)
 
       expect(agentCR.spec.tools).toBeUndefined()
     })
+
+    test('should handle tools with custom description', () => {
+      const requestWithDescription = {
+        ...mockAgentRequest,
+        spec: {
+          ...mockAgentRequest.spec,
+          tools: [
+            {
+              type: 'knowledgeBase',
+              name: 'test-kb',
+              description: 'Custom description for the knowledge base',
+            },
+          ],
+        },
+      }
+
+      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithDescription)
+
+      expect(agentCR.spec.tools).toEqual([
+        {
+          type: 'knowledgeBase',
+          name: 'test-kb',
+          description: 'Custom description for the knowledge base',
+          endpoint: undefined,
+        },
+      ])
+    })
   })
 
   describe('toRecord', () => {
@@ -113,7 +148,15 @@ describe('AkamaiAgentCR', () => {
         spec: {
           foundationModel: 'gpt-4',
           agentInstructions: 'You are a helpful assistant',
-          knowledgeBase: 'test-kb',
+          tools: [
+            {
+              type: 'knowledgeBase',
+              name: 'test-kb',
+              description:
+                'Search the test-kb knowledge base for relevant information. Use this when you need factual information, documentation, or specific details stored in the knowledge base.',
+              endpoint: undefined,
+            },
+          ],
         },
         status: {
           conditions: [
@@ -128,19 +171,48 @@ describe('AkamaiAgentCR', () => {
       })
     })
 
-    test('should handle empty knowledgeBase in response', () => {
-      const requestWithoutKB = {
+    test('should handle empty tools array in response', () => {
+      const requestWithoutTools = {
         ...mockAgentRequest,
         spec: {
           ...mockAgentRequest.spec,
-          knowledgeBase: undefined,
+          tools: undefined,
         },
       }
 
-      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutKB)
+      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutTools)
       const response = agentCR.toApiResponse('team-123')
 
-      expect(response.spec.knowledgeBase).toBe('')
+      expect(response.spec.tools).toBeUndefined()
+    })
+
+    test('should preserve custom description and endpoint in response', () => {
+      const requestWithDetails = {
+        ...mockAgentRequest,
+        spec: {
+          ...mockAgentRequest.spec,
+          tools: [
+            {
+              type: 'knowledgeBase',
+              name: 'test-kb',
+              description: 'Custom KB description',
+              endpoint: 'https://api.example.com/kb',
+            },
+          ],
+        },
+      }
+
+      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithDetails)
+      const response = agentCR.toApiResponse('team-123')
+
+      expect(response.spec.tools).toEqual([
+        {
+          type: 'knowledgeBase',
+          name: 'test-kb',
+          description: 'Custom KB description',
+          endpoint: 'https://api.example.com/kb',
+        },
+      ])
     })
   })
 
@@ -207,7 +279,7 @@ describe('AkamaiAgentCR', () => {
         apiVersion: 'akamai.com/v1',
         kind: 'Agent',
         metadata: { name: 'existing-agent', namespace: 'team-456' },
-        spec: { foundationModel: 'gpt-3.5', systemPrompt: 'Test prompt' },
+        spec: { foundationModel: 'gpt-3.5', agentInstructions: 'Test prompt' },
       }
 
       const result = AkamaiAgentCR.fromCR(crObject)
diff --git a/src/ai/AkamaiAgentCR.ts b/src/ai/AkamaiAgentCR.ts
@@ -18,7 +18,7 @@ export class AkamaiAgentCR {
   }
   public spec: {
     foundationModel: string
-    systemPrompt: string
+    agentInstructions: string
     tools?: Array<{
       type: string
       name: string
@@ -42,16 +42,17 @@ export class AkamaiAgentCR {
     }
     this.spec = {
       foundationModel: request.spec.foundationModel,
-      systemPrompt: request.spec.agentInstructions,
-      tools: request.spec.knowledgeBase
-        ? [
-            {
-              type: 'knowledgeBase',
-              name: request.spec.knowledgeBase,
-              description: `Search the ${request.spec.knowledgeBase} knowledge base for relevant information. Use this when you need factual information, documentation, or specific details stored in the knowledge base.`,
-            },
-          ]
-        : undefined,
+      agentInstructions: request.spec.agentInstructions,
+      tools: request.spec.tools?.map((tool) => ({
+        type: tool.type,
+        name: tool.name,
+        description:
+          tool.description ||
+          (tool.type === 'knowledgeBase'
+            ? `Search the ${tool.name} knowledge base for relevant information. Use this when you need factual information, documentation, or specific details stored in the knowledge base.`
+            : undefined),
+        endpoint: tool.endpoint,
+      })),
     }
   }
 
@@ -67,9 +68,6 @@ export class AkamaiAgentCR {
 
   // Transform to API response format
   toApiResponse(teamId: string): AplAgentResponse {
-    // Extract knowledgeBase from tools array (find first knowledgeBase tool)
-    const knowledgeBaseTool = this.spec.tools?.find((tool) => tool.type === 'knowledgeBase')
-
     return {
       kind: 'AkamaiAgent',
       metadata: {
@@ -81,8 +79,13 @@ export class AkamaiAgentCR {
       },
       spec: {
         foundationModel: this.spec.foundationModel,
-        agentInstructions: this.spec.systemPrompt,
-        knowledgeBase: knowledgeBaseTool?.name || '',
+        agentInstructions: this.spec.agentInstructions,
+        tools: this.spec.tools?.map((tool) => ({
+          type: tool.type,
+          name: tool.name,
+          description: tool.description,
+          endpoint: tool.endpoint,
+        })),
       },
       status: {
         conditions: [
diff --git a/src/ai/aiModelHandler.test.ts b/src/ai/aiModelHandler.test.ts
@@ -52,10 +52,10 @@ describe('aiModelHandler', () => {
       expect(result).toEqual({
         kind: 'AplAIModel',
         metadata: {
-          name: 'gpt-4-deployment',
+          name: 'gpt-4',
         },
         spec: {
-          displayName: 'gpt-4-deployment',
+          displayName: 'gpt-4',
           modelEndpoint: 'http://gpt-4-deployment.ai-models.svc.cluster.local',
           modelType: 'foundation',
           modelDimension: 1536,
@@ -97,8 +97,8 @@ describe('aiModelHandler', () => {
 
       const result = transformK8sDeploymentToAplAIModel(deploymentWithModelName)
 
-      expect(result.metadata.name).toBe('some-deployment-name')
-      expect(result.spec.displayName).toBe('some-deployment-name')
+      expect(result.metadata.name).toBe('custom-model-name')
+      expect(result.spec.displayName).toBe('custom-model-name')
     })
 
     test('should use modelName from labels when deployment name is missing', () => {
@@ -231,7 +231,7 @@ describe('aiModelHandler', () => {
 
       const result = transformK8sDeploymentToAplAIModel(deploymentWithoutMetadata)
 
-      expect(result.metadata.name).toBeUndefined()
+      expect(result.metadata.name).toBe('')
       expect(result.spec.modelEndpoint).toBe('http://undefined.undefined.svc.cluster.local')
     })
   })
@@ -244,7 +244,7 @@ describe('aiModelHandler', () => {
 
       expect(result).toHaveLength(1)
       expect(result[0].kind).toBe('AplAIModel')
-      expect(result[0].metadata.name).toBe('gpt-4-deployment')
+      expect(result[0].metadata.name).toBe('gpt-4')
       expect(mockedGetDeploymentsWithAIModelLabels).toHaveBeenCalledTimes(1)
     })
 
@@ -276,8 +276,8 @@ describe('aiModelHandler', () => {
       const result = await getAIModels()
 
       expect(result).toHaveLength(2)
-      expect(result[0].metadata.name).toBe('gpt-4-deployment')
-      expect(result[1].metadata.name).toBe('embedding-model')
+      expect(result[0].metadata.name).toBe('gpt-4')
+      expect(result[1].metadata.name).toBe('text-embedding-ada-002')
     })
 
     test('should propagate errors from k8s module', async () => {
diff --git a/src/k8s_operations.test.ts b/src/k8s_operations.test.ts
@@ -1,5 +1,23 @@
 import { CoreV1Api } from '@kubernetes/client-node'
 import { getCloudttyActiveTime, getLogTime } from './k8s_operations'
+
+// Mock the KubeConfig
+jest.mock('@kubernetes/client-node', () => {
+  const actual = jest.requireActual('@kubernetes/client-node')
+  return {
+    ...actual,
+    KubeConfig: jest.fn().mockImplementation(() => ({
+      loadFromDefault: jest.fn(),
+      makeApiClient: jest.fn((apiClientType) => {
+        if (apiClientType === actual.CoreV1Api) {
+          return new actual.CoreV1Api()
+        }
+        return {}
+      }),
+    })),
+  }
+})
+
 describe('getCloudttyLogTime', () => {
   test('should return the timestamp for a valid log timestamp', () => {
     const timestampMatch = ['[2023/10/10 00:00:00:0000]', '2023/10/10 00:00:00:0000']
@@ -17,7 +35,7 @@ describe('getCloudttyLogTime', () => {
 
 describe('getCloudttyActiveTime', () => {
   afterEach(() => {
-    jest.restoreAllMocks()
+    jest.clearAllMocks()
   })
 
   test('should return the time difference if no clients', async () => {
diff --git a/src/openapi/agent.yaml b/src/openapi/agent.yaml
@@ -19,10 +19,6 @@ Agent:
 AplAgentSpec:
   type: object
   properties:
-    knowledgeBase:
-      type: string
-      description: Name of the knowledge base to use
-      example: "company-docs"
     foundationModel:
       type: string
       description: Name of the foundation model
@@ -31,6 +27,30 @@ AplAgentSpec:
       type: string
       description: Custom instructions for the agent
       example: "You are a helpful assistant that provides concise answers."
+    tools:
+      type: array
+      description: Tools available to the agent
+      items:
+        type: object
+        properties:
+          type:
+            type: string
+            description: Type of the tool
+            example: "knowledgeBase"
+          name:
+            type: string
+            description: Name of the tool resource
+            example: "company-docs"
+          description:
+            type: string
+            description: Description of what the tool does
+            example: "Search the company-docs knowledge base for relevant information"
+          endpoint:
+            type: string
+            description: Optional endpoint URL for the tool
+        required:
+          - type
+          - name
   required:
     - foundationModel
     - agentInstructions
diff --git a/src/otomi-stack.ts b/src/otomi-stack.ts
@@ -2488,7 +2488,7 @@ export default class OtomiStack {
           spec: {
             foundationModel: data.spec?.foundationModel ?? existingAgent.spec.foundationModel,
             agentInstructions: data.spec?.agentInstructions ?? existingAgent.spec.agentInstructions,
-            knowledgeBase: data.spec?.knowledgeBase ?? existingAgent.spec.knowledgeBase,
+            tools: (data.spec?.tools ?? existingAgent.spec.tools) as typeof existingAgent.spec.tools,
           },
         })
 
diff --git a/src/services/TeamConfigService.test.ts b/src/services/TeamConfigService.test.ts
diff --git a/src/services/TeamConfigService.ts b/src/services/TeamConfigService.ts
