Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 3 additions & 6 deletions src/ai/AkamaiAgentCR.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@ describe('AkamaiAgentCR', () => {
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,
},
])
})
Expand Down Expand Up @@ -111,7 +110,6 @@ describe('AkamaiAgentCR', () => {
type: 'knowledgeBase',
name: 'test-kb',
description: 'Custom description for the knowledge base',
endpoint: undefined,
},
])
})
Expand Down Expand Up @@ -154,7 +152,6 @@ describe('AkamaiAgentCR', () => {
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,
},
],
},
Expand Down Expand Up @@ -186,7 +183,7 @@ describe('AkamaiAgentCR', () => {
expect(response.spec.tools).toBeUndefined()
})

test('should preserve custom description and endpoint in response', () => {
test('should preserve custom description and apiUrl in response', () => {
const requestWithDetails = {
...mockAgentRequest,
spec: {
Expand All @@ -196,7 +193,7 @@ describe('AkamaiAgentCR', () => {
type: 'knowledgeBase',
name: 'test-kb',
description: 'Custom KB description',
endpoint: 'https://api.example.com/kb',
apiUrl: 'https://api.example.com/kb',
},
],
},
Expand All @@ -210,7 +207,7 @@ describe('AkamaiAgentCR', () => {
type: 'knowledgeBase',
name: 'test-kb',
description: 'Custom KB description',
endpoint: 'https://api.example.com/kb',
apiUrl: 'https://api.example.com/kb',
},
])
})
Expand Down
81 changes: 61 additions & 20 deletions src/ai/AkamaiAgentCR.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,20 @@ export class AkamaiAgentCR {
}
public spec: {
foundationModel: string
foundationModelEndpoint?: string
agentInstructions: string
routes?: Array<{
agent: string
condition: string
apiUrl: string
apiKey?: string
}>
tools?: Array<{
type: string
name: string
description?: string
endpoint?: string
apiUrl?: string
apiKey?: string
}>
}

Expand All @@ -42,17 +50,29 @@ export class AkamaiAgentCR {
}
this.spec = {
foundationModel: request.spec.foundationModel,
...(request.spec.foundationModelEndpoint && { foundationModelEndpoint: request.spec.foundationModelEndpoint }),
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,
})),
...(request.spec.routes && {
routes: request.spec.routes.map((route) => ({
agent: route.agent,
condition: route.condition,
apiUrl: route.apiUrl,
...(route.apiKey && { apiKey: route.apiKey }),
})),
}),
...(request.spec.tools && {
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),
...(tool.apiUrl && { apiUrl: tool.apiUrl }),
...(tool.apiKey && { apiKey: tool.apiKey }),
})),
}),
}
}

Expand All @@ -79,13 +99,25 @@ export class AkamaiAgentCR {
},
spec: {
foundationModel: this.spec.foundationModel,
...(this.spec.foundationModelEndpoint && { foundationModelEndpoint: this.spec.foundationModelEndpoint }),
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 }),
})),
...(this.spec.routes && {
routes: this.spec.routes.map((route) => ({
agent: route.agent,
condition: route.condition,
apiUrl: route.apiUrl,
...(route.apiKey && { apiKey: route.apiKey }),
})),
}),
...(this.spec.tools && {
tools: this.spec.tools.map((tool) => ({
type: tool.type,
name: tool.name,
...(tool.description && { description: tool.description }),
...(tool.apiUrl && { apiUrl: tool.apiUrl }),
...(tool.apiKey && { apiKey: tool.apiKey }),
})),
}),
},
status: {
conditions: [
Expand All @@ -103,15 +135,24 @@ export class AkamaiAgentCR {
// Static factory method
static async create(teamId: string, agentName: string, request: AplAgentRequest): Promise<AkamaiAgentCR> {
const aiModels = await getAIModels()
const embeddingModel = aiModels.find(
const foundationModel = aiModels.find(
(model) => model.metadata.name === request.spec.foundationModel && model.spec.modelType === 'foundation',
)

if (!embeddingModel) {
if (!foundationModel) {
throw new K8sResourceNotFound('Foundation model', `Foundation model '${request.spec.foundationModel}' not found`)
}

return new AkamaiAgentCR(teamId, agentName, request)
// Create enriched request with foundationModelEndpoint from the model
const enrichedRequest: AplAgentRequest = {
...request,
spec: {
...request.spec,
foundationModelEndpoint: foundationModel.spec.modelEndpoint,
},
}

return new AkamaiAgentCR(teamId, agentName, enrichedRequest)
}

// Static method to create from existing CR (for transformation)
Expand Down
52 changes: 31 additions & 21 deletions src/ai/aiModelHandler.test.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,15 @@
import { V1Deployment } from '@kubernetes/client-node'
import { getAIModels, transformK8sDeploymentToAplAIModel } from './aiModelHandler'
import { getAIModels, transformK8sWorkloadToAplAIModel } from './aiModelHandler'
import * as k8s from './k8s'

// Mock the k8s module
jest.mock('./k8s')
const mockedGetDeploymentsWithAIModelLabels = k8s.getDeploymentsWithAIModelLabels as jest.MockedFunction<
typeof k8s.getDeploymentsWithAIModelLabels
>
const mockedGetStatefulSetsWithAIModelLabels = k8s.getStatefulSetsWithAIModelLabels as jest.MockedFunction<
typeof k8s.getStatefulSetsWithAIModelLabels
>

describe('aiModelHandler', () => {
const mockDeployment: V1Deployment = {
Expand All @@ -19,6 +22,7 @@ describe('aiModelHandler', () => {
modelNameTitle: 'GPT-4o-mini',
modelType: 'foundation',
modelDimension: '1536',
'serving.knative.dev/service': 'gpt-4-deployment',
},
},
status: {
Expand Down Expand Up @@ -47,18 +51,18 @@ describe('aiModelHandler', () => {
jest.clearAllMocks()
})

describe('transformK8sDeploymentToAplAIModel', () => {
describe('transformK8sWorkloadToAplAIModel', () => {
test('should transform K8s deployment to AplAIModel with all fields', () => {
const result = transformK8sDeploymentToAplAIModel(mockDeployment)
const result = transformK8sWorkloadToAplAIModel(mockDeployment)

expect(result).toEqual({
kind: 'AplAIModel',
metadata: {
name: 'gpt-4',
},
spec: {
displayName: 'GPT-4o-mini',
modelEndpoint: 'http://gpt-4.ai-models.svc.cluster.local/openai/v1',
displayName: 'gpt-4',
modelEndpoint: 'http://gpt-4-deployment.ai-models.svc.cluster.local/openai/v1',
modelType: 'foundation',
modelDimension: 1536,
},
Expand Down Expand Up @@ -97,10 +101,10 @@ describe('aiModelHandler', () => {
},
}

const result = transformK8sDeploymentToAplAIModel(deploymentWithModelName)
const result = transformK8sWorkloadToAplAIModel(deploymentWithModelName)

expect(result.metadata.name).toBe('custom-model-name')
expect(result.spec.displayName).toBe('GPT-4o-mini')
expect(result.spec.displayName).toBe('custom-model-name')
})

test('should use modelName from labels when deployment name is missing', () => {
Expand All @@ -116,10 +120,10 @@ describe('aiModelHandler', () => {
},
}

const result = transformK8sDeploymentToAplAIModel(deploymentWithoutName)
const result = transformK8sWorkloadToAplAIModel(deploymentWithoutName)

expect(result.metadata.name).toBe('custom-model-name')
expect(result.spec.displayName).toBe('GPT-4o-mini')
expect(result.spec.displayName).toBe('custom-model-name')
})

test('should handle deployment without labels', () => {
Expand All @@ -131,7 +135,7 @@ describe('aiModelHandler', () => {
},
}

const result = transformK8sDeploymentToAplAIModel(deploymentWithoutLabels)
const result = transformK8sWorkloadToAplAIModel(deploymentWithoutLabels)

expect(result.metadata.name).toBe('test-deployment')
expect(result.spec.modelType).toBeUndefined()
Expand All @@ -150,7 +154,7 @@ describe('aiModelHandler', () => {
},
}

const result = transformK8sDeploymentToAplAIModel(deploymentWithoutDimension)
const result = transformK8sWorkloadToAplAIModel(deploymentWithoutDimension)

expect(result.spec.modelDimension).toBeUndefined()
})
Expand All @@ -164,9 +168,9 @@ describe('aiModelHandler', () => {
},
}

const result = transformK8sDeploymentToAplAIModel(deploymentWithoutNamespace)
const result = transformK8sWorkloadToAplAIModel(deploymentWithoutNamespace)

expect(result.spec.modelEndpoint).toBe('http://gpt-4.undefined.svc.cluster.local/openai/v1')
expect(result.spec.modelEndpoint).toBe('http://gpt-4-deployment.undefined.svc.cluster.local/openai/v1')
})

test('should handle deployment without status conditions', () => {
Expand All @@ -178,7 +182,7 @@ describe('aiModelHandler', () => {
},
}

const result = transformK8sDeploymentToAplAIModel(deploymentWithoutConditions)
const result = transformK8sWorkloadToAplAIModel(deploymentWithoutConditions)

expect(result.status.conditions).toEqual([])
expect(result.status.phase).toBe('NotReady')
Expand All @@ -193,13 +197,13 @@ describe('aiModelHandler', () => {
},
}

const result = transformK8sDeploymentToAplAIModel(notReadyDeployment)
const result = transformK8sWorkloadToAplAIModel(notReadyDeployment)

expect(result.status.phase).toBe('NotReady')
})

test('should set phase to Ready when has ready replicas', () => {
const result = transformK8sDeploymentToAplAIModel(mockDeployment)
const result = transformK8sWorkloadToAplAIModel(mockDeployment)

expect(result.status.phase).toBe('Ready')
})
Expand All @@ -221,7 +225,7 @@ describe('aiModelHandler', () => {
},
}

const result = transformK8sDeploymentToAplAIModel(deploymentWithFalseCondition)
const result = transformK8sWorkloadToAplAIModel(deploymentWithFalseCondition)

expect(result.status.conditions?.[0]?.status).toBe(false)
})
Expand All @@ -231,35 +235,39 @@ describe('aiModelHandler', () => {
status: mockDeployment.status,
} as V1Deployment

const result = transformK8sDeploymentToAplAIModel(deploymentWithoutMetadata)
const result = transformK8sWorkloadToAplAIModel(deploymentWithoutMetadata)

expect(result.metadata.name).toBe('')
expect(result.spec.modelEndpoint).toBe('http://.undefined.svc.cluster.local/openai/v1')
expect(result.spec.modelEndpoint).toBe('http://.undefined.svc.cluster.local/v1')
})
})

describe('getAIModels', () => {
test('should return transformed AI models from deployments', async () => {
mockedGetDeploymentsWithAIModelLabels.mockResolvedValue([mockDeployment])
mockedGetStatefulSetsWithAIModelLabels.mockResolvedValue([])

const result = await getAIModels()

expect(result).toHaveLength(1)
expect(result[0].kind).toBe('AplAIModel')
expect(result[0].metadata.name).toBe('gpt-4')
expect(mockedGetDeploymentsWithAIModelLabels).toHaveBeenCalledTimes(1)
expect(mockedGetStatefulSetsWithAIModelLabels).toHaveBeenCalledTimes(1)
})

test('should return empty array when no deployments found', async () => {
test('should return empty array when no deployments or statefulsets found', async () => {
mockedGetDeploymentsWithAIModelLabels.mockResolvedValue([])
mockedGetStatefulSetsWithAIModelLabels.mockResolvedValue([])

const result = await getAIModels()

expect(result).toEqual([])
expect(mockedGetDeploymentsWithAIModelLabels).toHaveBeenCalledTimes(1)
expect(mockedGetStatefulSetsWithAIModelLabels).toHaveBeenCalledTimes(1)
})

test('should handle multiple deployments', async () => {
test('should handle multiple deployments and statefulsets', async () => {
const secondDeployment = {
...mockDeployment,
metadata: {
Expand All @@ -274,6 +282,7 @@ describe('aiModelHandler', () => {
}

mockedGetDeploymentsWithAIModelLabels.mockResolvedValue([mockDeployment, secondDeployment])
mockedGetStatefulSetsWithAIModelLabels.mockResolvedValue([])

const result = await getAIModels()

Expand All @@ -285,6 +294,7 @@ describe('aiModelHandler', () => {
test('should propagate errors from k8s module', async () => {
const error = new Error('K8s API error')
mockedGetDeploymentsWithAIModelLabels.mockRejectedValue(error)
mockedGetStatefulSetsWithAIModelLabels.mockResolvedValue([])

await expect(getAIModels()).rejects.toThrow('K8s API error')
})
Expand Down
Loading
Loading