feat: adjust kb custom resource

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,8 +53,16 @@ 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.knowledgeBase).toBe('test-kb')
+      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,
+        },
+      ])
     })
 
     test('should set teamId label and not merge custom labels', () => {
@@ -62,18 +75,45 @@ 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,
+          tools: undefined,
+        },
+      }
+
+      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,
-          knowledgeBase: undefined,
+          tools: [
+            {
+              type: 'knowledgeBase',
+              name: 'test-kb',
+              description: 'Custom description for the knowledge base',
+            },
+          ],
         },
       }
 
-      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutKB)
+      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithDescription)
 
-      expect(agentCR.spec.knowledgeBase).toBeUndefined()
+      expect(agentCR.spec.tools).toEqual([
+        {
+          type: 'knowledgeBase',
+          name: 'test-kb',
+          description: 'Custom description for the knowledge base',
+          endpoint: undefined,
+        },
+      ])
     })
   })
 
@@ -108,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: [
@@ -123,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,
+          tools: undefined,
+        },
+      }
+
+      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutTools)
+      const response = agentCR.toApiResponse('team-123')
+
+      expect(response.spec.tools).toBeUndefined()
+    })
+
+    test('should preserve custom description and endpoint in response', () => {
+      const requestWithDetails = {
         ...mockAgentRequest,
         spec: {
           ...mockAgentRequest.spec,
-          knowledgeBase: undefined,
+          tools: [
+            {
+              type: 'knowledgeBase',
+              name: 'test-kb',
+              description: 'Custom KB description',
+              endpoint: 'https://api.example.com/kb',
+            },
+          ],
         },
       }
 
-      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutKB)
+      const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithDetails)
       const response = agentCR.toApiResponse('team-123')
 
-      expect(response.spec.knowledgeBase).toBe('')
+      expect(response.spec.tools).toEqual([
+        {
+          type: 'knowledgeBase',
+          name: 'test-kb',
+          description: 'Custom KB description',
+          endpoint: 'https://api.example.com/kb',
+        },
+      ])
     })
   })
 
@@ -202,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)

src/ai/AkamaiAgentCR.ts

@@ -18,8 +18,13 @@ export class AkamaiAgentCR {
   }
   public spec: {
     foundationModel: string
-    systemPrompt: string
-    knowledgeBase?: string
+    agentInstructions: string
+    tools?: Array<{
+      type: string
+      name: string
+      description?: string
+      endpoint?: string
+    }>
   }
 
   constructor(teamId: string, agentName: string, request: AplAgentRequest) {
@@ -37,8 +42,17 @@ export class AkamaiAgentCR {
     }
     this.spec = {
       foundationModel: request.spec.foundationModel,
-      systemPrompt: request.spec.agentInstructions,
-      knowledgeBase: request.spec.knowledgeBase,
+      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,
+      })),
     }
   }
 
@@ -65,8 +79,13 @@ export class AkamaiAgentCR {
       },
       spec: {
         foundationModel: this.spec.foundationModel,
-        agentInstructions: this.spec.systemPrompt,
-        knowledgeBase: this.spec.knowledgeBase || '',
+        agentInstructions: this.spec.agentInstructions,
+        tools: this.spec.tools?.map((tool) => ({
+          type: tool.type,
+          name: tool.name,
+          ...(tool.description && { description: tool.description }),
+          ...(tool.endpoint && { endpoint: tool.endpoint }),
+        })),
       },
       status: {
         conditions: [

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 () => {

src/ai/aiModelHandler.ts

@@ -14,7 +14,7 @@ function getConditions(deployment: V1Deployment) {
 
 export function transformK8sDeploymentToAplAIModel(deployment: V1Deployment): AplAIModelResponse {
   const labels = deployment.metadata?.labels || {}
-  const modelName = deployment.metadata?.name || labels.modelName
+  const modelName = labels.modelName || deployment.metadata?.name || ''
 
   // Convert K8s deployment conditions to schema format
   const conditions = getConditions(deployment)

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 () => {

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

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,
           },
         })