Skip to content

Commit 1b4cab3

Browse files
CasLubbersmerllsvcAPLBot
authored
feat: add knowledge base apis (#802)
* feat: added API spec for agent inference platform * feat: add agent crud api and fix comments * feat: add ai model get api * feat: add authz tests * feat: add knowledgebase apis * feat: update prettier * feat: create classes for CR * feat: rename dir to knowledgebases * feat: add agent api * feat: add agent api * feat: add validation and tests * feat: update pipelineParameters to snake_case * fix: api authz test --------- Co-authored-by: Matthias Erll <merll@akamai.com> Co-authored-by: svcAPLBot <174728082+svcAPLBot@users.noreply.github.com>
1 parent d86dfb4 commit 1b4cab3

26 files changed

+2388
-17
lines changed

src/ai/AkamaiAgentCR.test.ts

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
import { AkamaiAgentCR } from './AkamaiAgentCR'
2+
import { AplAgentRequest } from 'src/otomi-models'
3+
import { K8sResourceNotFound } from '../error'
4+
import * as aiModelHandler from './aiModelHandler'
5+
6+
// Mock the aiModelHandler module
7+
jest.mock('./aiModelHandler')
8+
const mockedGetAIModels = aiModelHandler.getAIModels as jest.MockedFunction<typeof aiModelHandler.getAIModels>
9+
10+
describe('AkamaiAgentCR', () => {
11+
const mockFoundationModel = {
12+
kind: 'AplAIModel',
13+
metadata: { name: 'gpt-4' },
14+
spec: {
15+
displayName: 'GPT-4',
16+
modelEndpoint: 'http://gpt-4.ai.svc.cluster.local',
17+
modelType: 'foundation' as const,
18+
},
19+
status: {
20+
conditions: [],
21+
phase: 'Ready' as const,
22+
},
23+
}
24+
25+
const mockAgentRequest: AplAgentRequest = {
26+
kind: 'AkamaiAgent',
27+
metadata: {
28+
name: 'test-agent',
29+
},
30+
spec: {
31+
foundationModel: 'gpt-4',
32+
agentInstructions: 'You are a helpful assistant',
33+
knowledgeBase: 'test-kb',
34+
},
35+
}
36+
37+
beforeEach(() => {
38+
jest.clearAllMocks()
39+
})
40+
41+
describe('constructor', () => {
42+
test('should create AkamaiAgentCR with all properties', () => {
43+
const agentCR = new AkamaiAgentCR('team-123', 'test-agent', mockAgentRequest)
44+
45+
expect(agentCR.apiVersion).toBeDefined()
46+
expect(agentCR.kind).toBeDefined()
47+
expect(agentCR.metadata.name).toBe('test-agent')
48+
expect(agentCR.metadata.namespace).toBe('team-team-123')
49+
expect(agentCR.metadata.labels?.['apl.io/teamId']).toBe('team-123')
50+
expect(agentCR.spec.foundationModel).toBe('gpt-4')
51+
expect(agentCR.spec.systemPrompt).toBe('You are a helpful assistant')
52+
expect(agentCR.spec.knowledgeBase).toBe('test-kb')
53+
})
54+
55+
test('should set teamId label and not merge custom labels', () => {
56+
const agentCR = new AkamaiAgentCR('team-123', 'test-agent', mockAgentRequest)
57+
58+
expect(agentCR.metadata.labels).toEqual({
59+
'apl.io/teamId': 'team-123',
60+
})
61+
// Custom labels from request are not merged in constructor
62+
expect(agentCR.metadata.labels?.['custom-label']).toBeUndefined()
63+
})
64+
65+
test('should handle request without knowledgeBase', () => {
66+
const requestWithoutKB = {
67+
...mockAgentRequest,
68+
spec: {
69+
...mockAgentRequest.spec,
70+
knowledgeBase: undefined,
71+
},
72+
}
73+
74+
const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutKB)
75+
76+
expect(agentCR.spec.knowledgeBase).toBeUndefined()
77+
})
78+
})
79+
80+
describe('toRecord', () => {
81+
test('should return serializable record', () => {
82+
const agentCR = new AkamaiAgentCR('team-123', 'test-agent', mockAgentRequest)
83+
const record = agentCR.toRecord()
84+
85+
expect(record).toEqual({
86+
apiVersion: agentCR.apiVersion,
87+
kind: agentCR.kind,
88+
metadata: agentCR.metadata,
89+
spec: agentCR.spec,
90+
})
91+
})
92+
})
93+
94+
describe('toApiResponse', () => {
95+
test('should transform to API response format', () => {
96+
const agentCR = new AkamaiAgentCR('team-123', 'test-agent', mockAgentRequest)
97+
const response = agentCR.toApiResponse('team-123')
98+
99+
expect(response).toEqual({
100+
kind: 'AkamaiAgent',
101+
metadata: {
102+
name: 'test-agent',
103+
namespace: 'team-team-123',
104+
labels: {
105+
'apl.io/teamId': 'team-123',
106+
},
107+
},
108+
spec: {
109+
foundationModel: 'gpt-4',
110+
agentInstructions: 'You are a helpful assistant',
111+
knowledgeBase: 'test-kb',
112+
},
113+
status: {
114+
conditions: [
115+
{
116+
type: 'AgentDeployed',
117+
status: true,
118+
reason: 'Scheduled',
119+
message: 'Successfully deployed the Agent',
120+
},
121+
],
122+
},
123+
})
124+
})
125+
126+
test('should handle empty knowledgeBase in response', () => {
127+
const requestWithoutKB = {
128+
...mockAgentRequest,
129+
spec: {
130+
...mockAgentRequest.spec,
131+
knowledgeBase: undefined,
132+
},
133+
}
134+
135+
const agentCR = new AkamaiAgentCR('team-123', 'test-agent', requestWithoutKB)
136+
const response = agentCR.toApiResponse('team-123')
137+
138+
expect(response.spec.knowledgeBase).toBe('')
139+
})
140+
})
141+
142+
describe('create', () => {
143+
test('should create AkamaiAgentCR when foundation model exists', async () => {
144+
mockedGetAIModels.mockResolvedValue([mockFoundationModel as any])
145+
146+
const result = await AkamaiAgentCR.create('team-123', 'test-agent', mockAgentRequest)
147+
148+
expect(result).toBeInstanceOf(AkamaiAgentCR)
149+
expect(result.metadata.name).toBe('test-agent')
150+
expect(mockedGetAIModels).toHaveBeenCalledTimes(1)
151+
})
152+
153+
test('should throw K8sResourceNotFound when foundation model does not exist', async () => {
154+
mockedGetAIModels.mockResolvedValue([])
155+
156+
await expect(AkamaiAgentCR.create('team-123', 'test-agent', mockAgentRequest)).rejects.toThrow(
157+
K8sResourceNotFound,
158+
)
159+
await expect(AkamaiAgentCR.create('team-123', 'test-agent', mockAgentRequest)).rejects.toThrow(
160+
"Foundation model 'gpt-4' not found",
161+
)
162+
})
163+
164+
test('should throw K8sResourceNotFound when foundation model has wrong type', async () => {
165+
const embeddingModel = {
166+
...mockFoundationModel,
167+
spec: {
168+
...mockFoundationModel.spec,
169+
modelType: 'embedding' as const,
170+
},
171+
}
172+
mockedGetAIModels.mockResolvedValue([embeddingModel as any])
173+
174+
await expect(AkamaiAgentCR.create('team-123', 'test-agent', mockAgentRequest)).rejects.toThrow(
175+
K8sResourceNotFound,
176+
)
177+
})
178+
179+
test('should throw error when foundationModel is undefined', async () => {
180+
const requestWithoutModel = {
181+
...mockAgentRequest,
182+
spec: {
183+
...mockAgentRequest.spec,
184+
foundationModel: undefined,
185+
},
186+
}
187+
188+
mockedGetAIModels.mockResolvedValue([mockFoundationModel as any])
189+
190+
await expect(AkamaiAgentCR.create('team-123', 'test-agent', requestWithoutModel as any)).rejects.toThrow(
191+
K8sResourceNotFound,
192+
)
193+
await expect(AkamaiAgentCR.create('team-123', 'test-agent', requestWithoutModel as any)).rejects.toThrow(
194+
"Foundation model 'undefined' not found",
195+
)
196+
})
197+
})
198+
199+
describe('fromCR', () => {
200+
test('should create instance from existing CR object', () => {
201+
const crObject = {
202+
apiVersion: 'akamai.com/v1',
203+
kind: 'Agent',
204+
metadata: { name: 'existing-agent', namespace: 'team-456' },
205+
spec: { foundationModel: 'gpt-3.5', systemPrompt: 'Test prompt' },
206+
}
207+
208+
const result = AkamaiAgentCR.fromCR(crObject)
209+
210+
expect(result).toBeInstanceOf(AkamaiAgentCR)
211+
expect(result.metadata.name).toBe('existing-agent')
212+
expect(result.spec.foundationModel).toBe('gpt-3.5')
213+
})
214+
})
215+
})

src/ai/AkamaiAgentCR.ts

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
import { AplAgentRequest, AplAgentResponse } from 'src/otomi-models'
2+
import { AGENT_API_VERSION, AGENT_KIND, cleanEnv } from '../validators'
3+
import { K8sResourceNotFound } from '../error'
4+
import { getAIModels } from './aiModelHandler'
5+
6+
const env = cleanEnv({
7+
AGENT_API_VERSION,
8+
AGENT_KIND,
9+
})
10+
11+
export class AkamaiAgentCR {
12+
public apiVersion: string
13+
public kind: string
14+
public metadata: {
15+
name: string
16+
namespace: string
17+
labels?: Record<string, string>
18+
}
19+
public spec: {
20+
foundationModel: string
21+
systemPrompt: string
22+
knowledgeBase?: string
23+
}
24+
25+
constructor(teamId: string, agentName: string, request: AplAgentRequest) {
26+
const namespace = `team-${teamId}`
27+
28+
this.apiVersion = env.AGENT_API_VERSION
29+
this.kind = env.AGENT_KIND
30+
this.metadata = {
31+
...request.metadata,
32+
name: agentName,
33+
namespace,
34+
labels: {
35+
'apl.io/teamId': teamId,
36+
},
37+
}
38+
this.spec = {
39+
foundationModel: request.spec.foundationModel,
40+
systemPrompt: request.spec.agentInstructions,
41+
knowledgeBase: request.spec.knowledgeBase,
42+
}
43+
}
44+
45+
// Convert to plain object for serialization
46+
toRecord(): Record<string, any> {
47+
return {
48+
apiVersion: this.apiVersion,
49+
kind: this.kind,
50+
metadata: this.metadata,
51+
spec: this.spec,
52+
}
53+
}
54+
55+
// Transform to API response format
56+
toApiResponse(teamId: string): AplAgentResponse {
57+
return {
58+
kind: 'AkamaiAgent',
59+
metadata: {
60+
...this.metadata,
61+
labels: {
62+
'apl.io/teamId': teamId,
63+
...(this.metadata.labels || {}),
64+
},
65+
},
66+
spec: {
67+
foundationModel: this.spec.foundationModel,
68+
agentInstructions: this.spec.systemPrompt,
69+
knowledgeBase: this.spec.knowledgeBase || '',
70+
},
71+
status: {
72+
conditions: [
73+
{
74+
type: 'AgentDeployed',
75+
status: true,
76+
reason: 'Scheduled',
77+
message: 'Successfully deployed the Agent',
78+
},
79+
],
80+
},
81+
}
82+
}
83+
84+
// Static factory method
85+
static async create(teamId: string, agentName: string, request: AplAgentRequest): Promise<AkamaiAgentCR> {
86+
const aiModels = await getAIModels()
87+
const embeddingModel = aiModels.find(
88+
(model) => model.metadata.name === request.spec.foundationModel && model.spec.modelType === 'foundation',
89+
)
90+
91+
if (!embeddingModel) {
92+
throw new K8sResourceNotFound('Foundation model', `Foundation model '${request.spec.foundationModel}' not found`)
93+
}
94+
95+
return new AkamaiAgentCR(teamId, agentName, request)
96+
}
97+
98+
// Static method to create from existing CR (for transformation)
99+
static fromCR(cr: any): AkamaiAgentCR {
100+
const instance = Object.create(AkamaiAgentCR.prototype)
101+
return Object.assign(instance, cr)
102+
}
103+
}

0 commit comments

Comments
 (0)