Skip to content

Commit c6153ee

Browse files
tomkistomkis
authored
chore: rework beeai sdk (#1358)
* chore: rework beeai sdk Signed-off-by: Tomas Weiss <[email protected]> * chore: form extension Signed-off-by: Tomas Weiss <[email protected]> * chore: rework oauth for sdk Signed-off-by: Tomas Weiss <[email protected]> * fix: formatting Signed-off-by: Tomas Weiss <[email protected]> * fix: better naming Signed-off-by: Tomas Weiss <[email protected]> * chore: cleanup unneeded types Signed-off-by: Tomas Weiss <[email protected]> * chore: remove todo Signed-off-by: Tomas Weiss <[email protected]> * fix: proper naming Signed-off-by: Tomas Weiss <[email protected]> * fix: code review commennts Signed-off-by: Tomas Weiss <[email protected]> * fix: code review comments Signed-off-by: Tomas Weiss <[email protected]> --------- Signed-off-by: Tomas Weiss <[email protected]>
1 parent 425b4fb commit c6153ee

34 files changed

+467
-428
lines changed

‎apps/beeai-sdk-py/examples/secrets_agent.py‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,22 +22,26 @@ async def secrets_agent(
2222
input: Message,
2323
secrets: Annotated[
2424
SecretsExtensionServer,
25-
SecretsExtensionSpec.single_demand(name="Slack", description="Slack API key"),
25+
SecretsExtensionSpec(
26+
params=SecretsServiceExtensionParams(
27+
secret_demands={"ibm_cloud": SecretDemand(description="IBM Cloud API key", name="IBM Cloud")}
28+
)
29+
),
2630
],
2731
):
2832
"""Agent that uses request a secret that can be provided during runtime"""
2933
if secrets and secrets.data and secrets.data.secret_fulfillments:
30-
yield f"Slack API key: {secrets.data.secret_fulfillments['default'].secret}"
34+
yield f"IBM Cloud API key: {secrets.data.secret_fulfillments['ibm_cloud'].secret}"
3135
else:
3236
runtime_provided_secrets = await secrets.request_secrets(
3337
params=SecretsServiceExtensionParams(
34-
secret_demands={"default": SecretDemand(description="I really need Slack Key", name="Slack")}
38+
secret_demands={"ibm_cloud": SecretDemand(description="I really need IBM Cloud Key", name="IBM Cloud")}
3539
)
3640
)
3741
if runtime_provided_secrets and runtime_provided_secrets.secret_fulfillments:
38-
yield f"Slack API key: {runtime_provided_secrets.secret_fulfillments['default'].secret}"
42+
yield f"IBM Cloud API key: {runtime_provided_secrets.secret_fulfillments['ibm_cloud'].secret}"
3943
else:
40-
yield "No Slack API key provided"
44+
yield "No IBM Cloud API key provided"
4145

4246

4347
def run():

‎apps/beeai-sdk-ts/src/client/a2a/extensions/handle-agent-card.ts‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
/**
2+
* Copyright 2025 © BeeAI a Series of LF Projects, LLC
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import type { AgentCapabilities } from '@a2a-js/sdk';
7+
8+
import type { ContextToken } from '../../context/types';
9+
import type { EmbeddingDemands, EmbeddingFulfillments } from './services/embedding';
10+
import { embeddingExtension } from './services/embedding';
11+
import type { LLMDemands, LLMFulfillments } from './services/llm';
12+
import { llmExtension } from './services/llm';
13+
import type { MCPDemands, MCPFulfillments } from './services/mcp';
14+
import { mcpExtension } from './services/mcp';
15+
import type { OAuthDemands, OAuthFulfillments } from './services/oauth-provider';
16+
import { oauthProviderExtension } from './services/oauth-provider';
17+
import { platformApiExtension } from './services/platform';
18+
import type { SecretDemands, SecretFulfillments } from './services/secrets';
19+
import { secretsExtension } from './services/secrets';
20+
import type { FormDemands, FormFulfillments } from './ui/form';
21+
import { formExtension } from './ui/form';
22+
import { oauthRequestExtension } from './ui/oauth';
23+
import type { SettingsDemands, SettingsFulfillments } from './ui/settings';
24+
import { settingsExtension } from './ui/settings';
25+
import { extractServiceExtensionDemands, fulfillServiceExtensionDemand } from './utils';
26+
27+
export interface Fulfillments {
28+
llm: (demand: LLMDemands) => Promise<LLMFulfillments>;
29+
embedding: (demand: EmbeddingDemands) => Promise<EmbeddingFulfillments>;
30+
mcp: (demand: MCPDemands) => Promise<MCPFulfillments>;
31+
oauth: (demand: OAuthDemands) => Promise<OAuthFulfillments>;
32+
settings: (demand: SettingsDemands) => Promise<SettingsFulfillments>;
33+
secrets: (demand: SecretDemands) => Promise<SecretFulfillments>;
34+
form: (demand: FormDemands) => Promise<FormFulfillments | null>;
35+
oauthRedirectUri: () => string | null;
36+
getContextToken: () => ContextToken;
37+
}
38+
39+
const mcpExtensionExtractor = extractServiceExtensionDemands(mcpExtension);
40+
const llmExtensionExtractor = extractServiceExtensionDemands(llmExtension);
41+
const embeddingExtensionExtractor = extractServiceExtensionDemands(embeddingExtension);
42+
const oauthExtensionExtractor = extractServiceExtensionDemands(oauthProviderExtension);
43+
const settingsExtensionExtractor = extractServiceExtensionDemands(settingsExtension);
44+
const secretExtensionExtractor = extractServiceExtensionDemands(secretsExtension);
45+
const formExtensionExtractor = extractServiceExtensionDemands(formExtension);
46+
47+
const fulfillMcpDemand = fulfillServiceExtensionDemand(mcpExtension);
48+
const fulfillLlmDemand = fulfillServiceExtensionDemand(llmExtension);
49+
const fulfillEmbeddingDemand = fulfillServiceExtensionDemand(embeddingExtension);
50+
const fulfillOAuthDemand = fulfillServiceExtensionDemand(oauthProviderExtension);
51+
const fulfillSettingsDemand = fulfillServiceExtensionDemand(settingsExtension);
52+
const fulfillSecretDemand = fulfillServiceExtensionDemand(secretsExtension);
53+
const fulfillFormDemand = fulfillServiceExtensionDemand(formExtension);
54+
55+
export const handleAgentCard = (agentCard: { capabilities: AgentCapabilities }) => {
56+
const extensions = agentCard.capabilities.extensions ?? [];
57+
58+
const llmDemands = llmExtensionExtractor(extensions);
59+
const embeddingDemands = embeddingExtensionExtractor(extensions);
60+
const mcpDemands = mcpExtensionExtractor(extensions);
61+
const oauthDemands = oauthExtensionExtractor(extensions);
62+
const settingsDemands = settingsExtensionExtractor(extensions);
63+
const secretDemands = secretExtensionExtractor(extensions);
64+
const formDemands = formExtensionExtractor(extensions);
65+
66+
const resolveMetadata = async (fulfillments: Fulfillments) => {
67+
let fulfilledMetadata = {};
68+
69+
fulfilledMetadata = platformApiExtension(fulfilledMetadata, fulfillments.getContextToken());
70+
71+
if (llmDemands) {
72+
fulfilledMetadata = fulfillLlmDemand(fulfilledMetadata, await fulfillments.llm(llmDemands));
73+
}
74+
75+
if (embeddingDemands) {
76+
fulfilledMetadata = fulfillEmbeddingDemand(fulfilledMetadata, await fulfillments.embedding(embeddingDemands));
77+
}
78+
79+
if (mcpDemands) {
80+
fulfilledMetadata = fulfillMcpDemand(fulfilledMetadata, await fulfillments.mcp(mcpDemands));
81+
}
82+
83+
if (oauthDemands) {
84+
fulfilledMetadata = fulfillOAuthDemand(fulfilledMetadata, await fulfillments.oauth(oauthDemands));
85+
}
86+
87+
if (settingsDemands) {
88+
fulfilledMetadata = fulfillSettingsDemand(fulfilledMetadata, await fulfillments.settings(settingsDemands));
89+
}
90+
91+
if (secretDemands) {
92+
fulfilledMetadata = fulfillSecretDemand(fulfilledMetadata, await fulfillments.secrets(secretDemands));
93+
}
94+
95+
if (formDemands) {
96+
const formFulfillment = await fulfillments.form(formDemands);
97+
if (formFulfillment) {
98+
fulfilledMetadata = fulfillFormDemand(fulfilledMetadata, formFulfillment);
99+
}
100+
}
101+
102+
const oauthRedirectUri = fulfillments.oauthRedirectUri();
103+
if (oauthRedirectUri) {
104+
fulfilledMetadata = {
105+
...fulfilledMetadata,
106+
[oauthRequestExtension.getUri()]: {
107+
redirect_uri: oauthRedirectUri,
108+
},
109+
};
110+
}
111+
112+
return fulfilledMetadata;
113+
};
114+
115+
return {
116+
resolveMetadata,
117+
demands: { llmDemands, embeddingDemands, mcpDemands, oauthDemands, settingsDemands, secretDemands, formDemands },
118+
};
119+
};

‎apps/beeai-sdk-ts/src/client/a2a/extensions/handle-task-status-update.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
/**
2+
* Copyright 2025 © BeeAI a Series of LF Projects, LLC
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import type { TaskStatusUpdateEvent } from '@a2a-js/sdk';
7+
8+
import type { SecretDemands } from './services/secrets';
9+
import { secretsMessageExtension } from './services/secrets';
10+
import type { FormDemands } from './ui/form';
11+
import { formMessageExtension } from './ui/form';
12+
import { oauthRequestExtension } from './ui/oauth';
13+
import { extractUiExtensionData } from './utils';
14+
15+
const secretsMessageExtensionExtractor = extractUiExtensionData(secretsMessageExtension);
16+
const formMessageExtensionExtractor = extractUiExtensionData(formMessageExtension);
17+
const oauthRequestExtensionExtractor = extractUiExtensionData(oauthRequestExtension);
18+
19+
export enum TaskStatusUpdateType {
20+
SecretRequired = 'secret-required',
21+
FormRequired = 'form-required',
22+
OAuthRequired = 'oauth-required',
23+
}
24+
25+
export interface SecretRequiredResult {
26+
type: TaskStatusUpdateType.SecretRequired;
27+
demands: SecretDemands;
28+
}
29+
30+
export interface FormRequiredResult {
31+
type: TaskStatusUpdateType.FormRequired;
32+
form: FormDemands;
33+
}
34+
35+
export interface OAuthRequiredResult {
36+
type: TaskStatusUpdateType.OAuthRequired;
37+
url: string;
38+
}
39+
40+
export type TaskStatusUpdateResult = SecretRequiredResult | FormRequiredResult | OAuthRequiredResult;
41+
42+
export const handleTaskStatusUpdate = (event: TaskStatusUpdateEvent): TaskStatusUpdateResult[] => {
43+
const results: TaskStatusUpdateResult[] = [];
44+
45+
if (event.status.state === 'auth-required') {
46+
const secretRequired = secretsMessageExtensionExtractor(event.status.message?.metadata);
47+
const oauthRequired = oauthRequestExtensionExtractor(event.status.message?.metadata);
48+
49+
if (oauthRequired) {
50+
results.push({
51+
type: TaskStatusUpdateType.OAuthRequired,
52+
url: oauthRequired.authorization_endpoint_url,
53+
});
54+
}
55+
56+
if (secretRequired) {
57+
results.push({
58+
type: TaskStatusUpdateType.SecretRequired,
59+
demands: secretRequired,
60+
});
61+
}
62+
} else if (event.status.state === 'input-required') {
63+
const formRequired = formMessageExtensionExtractor(event.status.message?.metadata);
64+
if (formRequired) {
65+
results.push({
66+
type: TaskStatusUpdateType.FormRequired,
67+
form: formRequired,
68+
});
69+
}
70+
}
71+
72+
return results;
73+
};

‎apps/beeai-sdk-ts/src/client/a2a/extensions/services/platform.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ const getMetadata = (contextToken: ContextToken) => {
1414
};
1515
};
1616

17-
export const activePlatformExtension = (metadata: Record<string, unknown>, contextToken: ContextToken) => {
17+
export const platformApiExtension = (metadata: Record<string, unknown>, contextToken: ContextToken) => {
1818
return {
1919
...metadata,
2020
[URI]: getMetadata(contextToken),

‎apps/beeai-sdk-ts/src/client/a2a/extensions/ui/form.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -146,15 +146,15 @@ export type CheckboxField = z.infer<typeof checkboxField>;
146146

147147
export type FormField = z.infer<typeof fieldSchema>;
148148

149-
export type FormRender = z.infer<typeof renderSchema>;
150-
export type FormResponse = z.infer<typeof responseSchema>;
151-
export type FormResponseValue = FormResponse['values'][string];
149+
export type FormDemands = z.infer<typeof renderSchema>;
150+
export type FormFulfillments = z.infer<typeof responseSchema>;
151+
export type FormResponseValue = FormFulfillments['values'][string];
152152

153-
export const formMessageExtension: A2AUiExtension<typeof URI, FormRender> = {
153+
export const formMessageExtension: A2AUiExtension<typeof URI, FormDemands> = {
154154
getMessageMetadataSchema: () => z.object({ [URI]: renderSchema }).partial(),
155155
getUri: () => URI,
156156
};
157-
export const formExtension: A2AServiceExtension<typeof URI, z.infer<typeof renderSchema>, FormResponse> = {
157+
export const formExtension: A2AServiceExtension<typeof URI, z.infer<typeof renderSchema>, FormFulfillments> = {
158158
getDemandsSchema: () => renderSchema,
159159
getFulfillmentSchema: () => responseSchema,
160160
getUri: () => URI,

‎apps/beeai-sdk-ts/src/client/a2a/extensions/ui/settings.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,13 +70,13 @@ export type SettingsCheckboxGroupFieldValue = z.infer<typeof checkboxGroupFieldV
7070
export type SettingsSingleSelectFieldValue = z.infer<typeof singleSelectFieldValue>;
7171
export type SettingsFieldValue = z.infer<typeof settingsFieldValue>;
7272

73-
export type SettingsRender = z.infer<typeof settingsRenderSchema>;
74-
export type AgentRunSettings = z.infer<typeof agentRunSettingsSchema>;
73+
export type SettingsDemands = z.infer<typeof settingsRenderSchema>;
74+
export type SettingsFulfillments = z.infer<typeof agentRunSettingsSchema>;
7575

7676
export const settingsExtension: A2AServiceExtension<
7777
typeof URI,
7878
z.infer<typeof settingsRenderSchema>,
79-
AgentRunSettings
79+
SettingsFulfillments
8080
> = {
8181
getDemandsSchema: () => settingsRenderSchema,
8282
getFulfillmentSchema: () => agentRunSettingsSchema,

‎apps/beeai-sdk-ts/src/index.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,12 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6+
export { type Fulfillments, handleAgentCard } from './client/a2a/extensions/handle-agent-card';
7+
export {
8+
handleTaskStatusUpdate,
9+
type TaskStatusUpdateResult,
10+
TaskStatusUpdateType,
11+
} from './client/a2a/extensions/handle-task-status-update';
612
export * from './client/a2a/extensions/services/embedding';
713
export * from './client/a2a/extensions/services/llm';
814
export * from './client/a2a/extensions/services/mcp';

0 commit comments

Comments
 (0)