From ecd16bd1abeb01a9a79130d7a9ad5aed6ed8c245 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luk=C3=A1=C5=A1=20Jane=C4=8Dek?= Date: Fri, 31 Jan 2025 15:18:12 +0100 Subject: [PATCH] feat(projectSecrets): add CRUDL for project secrets (#40) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Lukáš Janeček Co-authored-by: Lukáš Janeček --- package.json | 2 +- pnpm-lock.yaml | 33 ++-- src/runs/dtos/run-submit-tool-inputs.ts | 66 +++++++ src/runs/dtos/run.ts | 17 ++ src/runs/entities/requiredAction.entity.ts | 6 +- src/runs/entities/requiredToolInput.entity.ts | 45 +++++ src/runs/entities/run.entity.ts | 3 +- .../execution/event-handlers/streaming.ts | 77 +------- src/runs/execution/tools/helpers.ts | 170 +++++++++++++++++- src/runs/runs.module.ts | 37 +++- src/runs/runs.service.ts | 79 +++++++- src/server.ts | 2 + src/tools/dtos/tool-create.ts | 3 +- src/tools/dtos/tool-secret-create.ts | 34 ++++ src/tools/dtos/tool-secret-delete.ts | 27 +++ src/tools/dtos/tool-secret-read.ts | 34 ++++ src/tools/dtos/tool-secret-update.ts | 38 ++++ src/tools/dtos/tool-secret.ts | 38 ++++ src/tools/dtos/tool-secrets-list.ts | 27 +++ src/tools/dtos/tool-update.ts | 1 + src/tools/entities/tool-secret.entity.ts | 58 ++++++ .../tool/code-interpreter-tool.entity.ts | 9 +- src/tools/tool-secrets.module.ts | 133 ++++++++++++++ src/tools/tool-secrets.service.ts | 117 ++++++++++++ src/tools/tools.service.ts | 2 + 25 files changed, 962 insertions(+), 96 deletions(-) create mode 100644 src/runs/dtos/run-submit-tool-inputs.ts create mode 100644 src/runs/entities/requiredToolInput.entity.ts create mode 100644 src/tools/dtos/tool-secret-create.ts create mode 100644 src/tools/dtos/tool-secret-delete.ts create mode 100644 src/tools/dtos/tool-secret-read.ts create mode 100644 src/tools/dtos/tool-secret-update.ts create mode 100644 src/tools/dtos/tool-secret.ts create mode 100644 src/tools/dtos/tool-secrets-list.ts create mode 100644 src/tools/entities/tool-secret.entity.ts create mode 100644 src/tools/tool-secrets.module.ts create mode 100644 src/tools/tool-secrets.service.ts diff --git a/package.json b/package.json index a6b6f5fc..219777b5 100644 --- a/package.json +++ b/package.json @@ -46,7 +46,7 @@ "@zilliz/milvus2-sdk-node": "^2.4.9", "ajv": "^8.17.1", "axios": "^1.7.7", - "bee-agent-framework": "0.0.57", + "bee-agent-framework": "0.0.58", "bullmq": "^5.34.6", "bullmq-otel": "^1.0.1", "cache-manager": "^5.7.6", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bb27d758..fea01d8d 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -40,10 +40,10 @@ importers: version: 6.2.9 '@mikro-orm/migrations-mongodb': specifier: 6.2.9 - version: 6.2.9(@mikro-orm/core@6.2.9)(@types/node@20.16.14)(gcp-metadata@6.1.0(encoding@0.1.13))(socks@2.8.3) + version: 6.2.9(@mikro-orm/core@6.2.9)(@types/node@20.16.14)(socks@2.8.3) '@mikro-orm/mongodb': specifier: 6.2.9 - version: 6.2.9(@mikro-orm/core@6.2.9)(gcp-metadata@6.1.0(encoding@0.1.13))(socks@2.8.3) + version: 6.2.9(@mikro-orm/core@6.2.9)(socks@2.8.3) '@mikro-orm/reflection': specifier: 6.2.9 version: 6.2.9(@mikro-orm/core@6.2.9) @@ -75,8 +75,8 @@ importers: specifier: ^1.7.7 version: 1.7.7 bee-agent-framework: - specifier: 0.0.57 - version: 0.0.57(@googleapis/customsearch@3.2.0(encoding@0.1.13))(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@zilliz/milvus2-sdk-node@2.5.3)(encoding@0.1.13)(google-auth-library@9.15.0(encoding@0.1.13))(ollama@0.5.12)(openai-chat-tokens@0.2.8)(openai@4.67.3(encoding@0.1.13)(zod@3.23.8))(yaml@2.7.0) + specifier: 0.0.58 + version: 0.0.58(@googleapis/customsearch@3.2.0(encoding@0.1.13))(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@zilliz/milvus2-sdk-node@2.5.3)(encoding@0.1.13)(google-auth-library@9.15.0(encoding@0.1.13))(ollama@0.5.12)(openai-chat-tokens@0.2.8)(openai@4.67.3(encoding@0.1.13)(zod@3.23.8))(yaml@2.7.0) bullmq: specifier: ^5.34.6 version: 5.34.6 @@ -2109,8 +2109,8 @@ packages: resolution: {integrity: sha512-4Bcg1P8xhUuqcii/S0Z9wiHIrQVPMermM1any+MX5GeGD7faD3/msQUDGLol9wOcz4/jbg/WJnGqoJF6LiBdtg==} engines: {node: '>=10.0.0'} - bee-agent-framework@0.0.57: - resolution: {integrity: sha512-Ne4ZIb/h2R2kWTXAW+FEe+xQrZQQHhooFoahE2gzxFXgJZP5v4XuuMEKFc2Epa534Mh6TCvazntmFT23qgmTiQ==} + bee-agent-framework@0.0.58: + resolution: {integrity: sha512-fOds1DmOFF09cePot6Tc1bcXnLOZhOCCQFolUseUvthclv+pNsLtrv7eZTI/rWAuQVq/X8P+QEo+3Cf0zFWKFg==} peerDependencies: '@aws-sdk/client-bedrock-runtime': ^3.687.0 '@elastic/elasticsearch': ^8.0.0 @@ -6162,12 +6162,12 @@ snapshots: - supports-color - tedious - '@mikro-orm/migrations-mongodb@6.2.9(@mikro-orm/core@6.2.9)(@types/node@20.16.14)(gcp-metadata@6.1.0(encoding@0.1.13))(socks@2.8.3)': + '@mikro-orm/migrations-mongodb@6.2.9(@mikro-orm/core@6.2.9)(@types/node@20.16.14)(socks@2.8.3)': dependencies: '@mikro-orm/core': 6.2.9 - '@mikro-orm/mongodb': 6.2.9(@mikro-orm/core@6.2.9)(gcp-metadata@6.1.0(encoding@0.1.13))(socks@2.8.3) + '@mikro-orm/mongodb': 6.2.9(@mikro-orm/core@6.2.9)(socks@2.8.3) fs-extra: 11.2.0 - mongodb: 6.7.0(gcp-metadata@6.1.0(encoding@0.1.13))(socks@2.8.3) + mongodb: 6.7.0(socks@2.8.3) umzug: 3.8.0(@types/node@20.16.14) transitivePeerDependencies: - '@aws-sdk/credential-providers' @@ -6179,11 +6179,11 @@ snapshots: - snappy - socks - '@mikro-orm/mongodb@6.2.9(@mikro-orm/core@6.2.9)(gcp-metadata@6.1.0(encoding@0.1.13))(socks@2.8.3)': + '@mikro-orm/mongodb@6.2.9(@mikro-orm/core@6.2.9)(socks@2.8.3)': dependencies: '@mikro-orm/core': 6.2.9 bson: 6.8.0 - mongodb: 6.7.0(gcp-metadata@6.1.0(encoding@0.1.13))(socks@2.8.3) + mongodb: 6.7.0(socks@2.8.3) transitivePeerDependencies: - '@aws-sdk/credential-providers' - '@mongodb-js/zstd' @@ -7838,7 +7838,7 @@ snapshots: basic-ftp@5.0.5: {} - bee-agent-framework@0.0.57(@googleapis/customsearch@3.2.0(encoding@0.1.13))(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@zilliz/milvus2-sdk-node@2.5.3)(encoding@0.1.13)(google-auth-library@9.15.0(encoding@0.1.13))(ollama@0.5.12)(openai-chat-tokens@0.2.8)(openai@4.67.3(encoding@0.1.13)(zod@3.23.8))(yaml@2.7.0): + bee-agent-framework@0.0.58(@googleapis/customsearch@3.2.0(encoding@0.1.13))(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@zilliz/milvus2-sdk-node@2.5.3)(encoding@0.1.13)(google-auth-library@9.15.0(encoding@0.1.13))(ollama@0.5.12)(openai-chat-tokens@0.2.8)(openai@4.67.3(encoding@0.1.13)(zod@3.23.8))(yaml@2.7.0): dependencies: '@ai-zen/node-fetch-event-source': 2.1.4(encoding@0.1.13) '@opentelemetry/api': 1.9.0 @@ -8652,7 +8652,7 @@ snapshots: debug: 4.3.6 enhanced-resolve: 5.17.1 eslint: 8.57.0 - eslint-module-utils: 2.8.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-plugin-import@2.29.1)(eslint@8.57.0))(eslint@8.57.0) + eslint-module-utils: 2.8.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1)(eslint@8.57.0) eslint-plugin-import: 2.29.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-typescript@3.6.1)(eslint@8.57.0) fast-glob: 3.3.2 get-tsconfig: 4.7.6 @@ -8664,7 +8664,7 @@ snapshots: - eslint-import-resolver-webpack - supports-color - eslint-module-utils@2.8.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-plugin-import@2.29.1)(eslint@8.57.0))(eslint@8.57.0): + eslint-module-utils@2.8.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1)(eslint@8.57.0): dependencies: debug: 3.2.7 optionalDependencies: @@ -8685,7 +8685,7 @@ snapshots: doctrine: 2.1.0 eslint: 8.57.0 eslint-import-resolver-node: 0.3.9 - eslint-module-utils: 2.8.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-plugin-import@2.29.1)(eslint@8.57.0))(eslint@8.57.0) + eslint-module-utils: 2.8.1(@typescript-eslint/parser@7.18.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1)(eslint@8.57.0) hasown: 2.0.2 is-core-module: 2.15.0 is-glob: 4.0.3 @@ -10006,13 +10006,12 @@ snapshots: '@types/whatwg-url': 11.0.5 whatwg-url: 13.0.0 - mongodb@6.7.0(gcp-metadata@6.1.0(encoding@0.1.13))(socks@2.8.3): + mongodb@6.7.0(socks@2.8.3): dependencies: '@mongodb-js/saslprep': 1.1.8 bson: 6.8.0 mongodb-connection-string-url: 3.0.1 optionalDependencies: - gcp-metadata: 6.1.0(encoding@0.1.13) socks: 2.8.3 ms@2.1.2: {} diff --git a/src/runs/dtos/run-submit-tool-inputs.ts b/src/runs/dtos/run-submit-tool-inputs.ts new file mode 100644 index 00000000..ae0b5291 --- /dev/null +++ b/src/runs/dtos/run-submit-tool-inputs.ts @@ -0,0 +1,66 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { FromSchema, JSONSchema } from 'json-schema-to-ts'; + +import { runParamsSchema, runSchema } from './run.js'; + +import { eventSchema } from '@/streaming/dtos/event.js'; + +export const runSubmitToolInputsParamsSchema = runParamsSchema; +export type RunSubmitToolInputsParams = FromSchema; + +export const runSubmitToolInputsBodySchema = { + type: 'object', + required: ['tool_inputs'], + additionalProperties: false, + properties: { + tool_inputs: { + type: 'array', + items: { + type: 'object', + additionalProperties: false, + required: ['tool_call_id', 'inputs'], + properties: { + tool_call_id: { type: 'string' }, + inputs: { + type: 'array', + items: { + type: 'object', + required: ['name', 'value'], + additionalProperties: false, + properties: { + name: { type: 'string' }, + value: { type: 'string' } + } + } + } + } + } + }, + stream: { + type: 'boolean', + nullable: true + } + } +} as const satisfies JSONSchema; +export type RunSubmitToolInputsBody = FromSchema; + +export const runSubmitToolInputsResponseSchema = runSchema; +export type RunSubmitToolInputsResponse = FromSchema; + +export const runSubmitToolInputsStreamSchema = eventSchema; +export type RunSubmitToolInputsStream = FromSchema; diff --git a/src/runs/dtos/run.ts b/src/runs/dtos/run.ts index c40d3e98..e23fe717 100644 --- a/src/runs/dtos/run.ts +++ b/src/runs/dtos/run.ts @@ -85,6 +85,23 @@ export const runSchema = { } } }, + { + required: ['type', 'submit_tool_inputs'], + properties: { + type: { const: 'submit_tool_inputs' }, + submit_tool_inputs: { + type: 'object', + required: ['tool_calls', 'input_fields'], + properties: { + tool_calls: { + type: 'array', + items: toolCallSchema + }, + input_fields: { type: 'array', items: { type: 'string' } } + } + } + } + }, { required: ['type', 'submit_tool_approvals'], properties: { diff --git a/src/runs/entities/requiredAction.entity.ts b/src/runs/entities/requiredAction.entity.ts index 489c9601..7e3b3e01 100644 --- a/src/runs/entities/requiredAction.entity.ts +++ b/src/runs/entities/requiredAction.entity.ts @@ -18,12 +18,14 @@ import { Embeddable, Enum, Property } from '@mikro-orm/core'; import { RequiredToolApprove } from './requiredToolApprove.entity'; import { RequiredToolOutput } from './requiredToolOutput.entity'; +import { RequiredToolInput } from './requiredToolInput.entity'; import { generatePrefixedObjectId } from '@/utils/id'; export enum RequiredActionType { OUTPUT = 'output', - APPROVE = 'approve' + APPROVE = 'approve', + INPUT = 'input' } @Embeddable({ abstract: true, discriminatorColumn: 'type' }) @@ -35,4 +37,4 @@ export abstract class RequiredAction { type!: RequiredActionType; } -export type AnyRequiredAction = RequiredToolApprove | RequiredToolOutput; +export type AnyRequiredAction = RequiredToolApprove | RequiredToolInput | RequiredToolOutput; diff --git a/src/runs/entities/requiredToolInput.entity.ts b/src/runs/entities/requiredToolInput.entity.ts new file mode 100644 index 00000000..78398788 --- /dev/null +++ b/src/runs/entities/requiredToolInput.entity.ts @@ -0,0 +1,45 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Embeddable, Embedded, Property } from '@mikro-orm/core'; + +import { RequiredAction, RequiredActionType } from './requiredAction.entity'; + +import { CodeInterpreterCall } from '@/tools/entities/tool-calls/code-interpreter-call.entity'; +import { FileSearchCall } from '@/tools/entities/tool-calls/file-search-call.entity'; +import { FunctionCall } from '@/tools/entities/tool-calls/function-call.entity'; +import { SystemCall } from '@/tools/entities/tool-calls/system-call.entity'; +import { UserCall } from '@/tools/entities/tool-calls/user-call.entity'; + +@Embeddable({ discriminatorValue: RequiredActionType.INPUT }) +export class RequiredToolInput extends RequiredAction { + type = RequiredActionType.INPUT; + + // Union must be defined in alphabetical order, otherwise Mikro-ORM won't discovered the auto-created virtual polymorphic entity + @Embedded({ object: true }) + toolCalls!: (CodeInterpreterCall | FileSearchCall | FunctionCall | SystemCall | UserCall)[]; + + @Property() + inputFields!: string[]; + + constructor({ toolCalls, inputFields }: RequiredToolInputInput) { + super(); + this.toolCalls = toolCalls; + this.inputFields = inputFields; + } +} + +export type RequiredToolInputInput = Pick; diff --git a/src/runs/entities/run.entity.ts b/src/runs/entities/run.entity.ts index 3dfe3dd7..81241e39 100644 --- a/src/runs/entities/run.entity.ts +++ b/src/runs/entities/run.entity.ts @@ -22,6 +22,7 @@ import { RUN_EXPIRATION_MILLISECONDS } from '../execution/constants.js'; import { RequiredToolApprove } from './requiredToolApprove.entity.js'; import { RequiredToolOutput } from './requiredToolOutput.entity.js'; import { ToolApproval } from './toolApproval.entity.js'; +import { RequiredToolInput } from './requiredToolInput.entity.js'; import { Assistant } from '@/assistants/assistant.entity.js'; import { Thread } from '@/threads/thread.entity.js'; @@ -75,7 +76,7 @@ export class Run extends PrincipalScopedEntity { // Union must be defined in alphabetical order, otherwise Mikro-ORM won't discovered the auto-created virtual polymorphic entity @Embedded({ object: true }) - requiredAction?: RequiredToolApprove | RequiredToolOutput; + requiredAction?: RequiredToolApprove | RequiredToolInput | RequiredToolOutput; @Embedded({ object: true }) toolApprovals?: ToolApproval[]; diff --git a/src/runs/execution/event-handlers/streaming.ts b/src/runs/execution/event-handlers/streaming.ts index 0a846d25..128103b4 100644 --- a/src/runs/execution/event-handlers/streaming.ts +++ b/src/runs/execution/event-handlers/streaming.ts @@ -20,7 +20,6 @@ import { ref } from '@mikro-orm/core'; import { Role } from 'bee-agent-framework/llms/primitives/message'; import { BeeCallbacks } from 'bee-agent-framework/agents/bee/types'; import { Summary } from 'prom-client'; -import { ToolError } from 'bee-agent-framework/tools/base'; import { StreamlitEvents as StreamlitEventsFramework, StreamlitRunOutput @@ -30,7 +29,12 @@ import { Agent } from '../constants'; import { AgentContext } from '@/runs/execution/execute.js'; import { getLogger } from '@/logger.js'; -import { createToolCall, finalizeToolCall } from '@/runs/execution/tools/helpers.js'; +import { + createToolCall, + finalizeToolCall, + requireToolApproval, + requireToolInput +} from '@/runs/execution/tools/helpers.js'; import { RunStep, RunStepStatus } from '@/run-steps/entities/run-step.entity.js'; import { RunStepToolCalls } from '@/run-steps/entities/details/run-step-tool-calls.entity.js'; import { ORM } from '@/database.js'; @@ -43,11 +47,6 @@ import { RunStatus } from '@/runs/entities/run.entity.js'; import { APIError } from '@/errors/error.entity.js'; import { jobRegistry } from '@/metrics.js'; import { EmitterEvent } from '@/run-steps/entities/emitter-event.entity'; -import { createApproveChannel, toRunDto } from '@/runs/runs.service'; -import { RequiredToolApprove } from '@/runs/entities/requiredToolApprove.entity'; -import { ToolApprovalType } from '@/runs/entities/toolApproval.entity'; -import { ToolType } from '@/tools/entities/tool/tool.entity'; -import { withRedisClient } from '@/redis.js'; import { Trace } from '@/observe/entities/trace.entity'; const agentToolExecutionTime = new Summary({ @@ -103,67 +102,9 @@ export function createBeeStreamingHandler(ctx: AgentContext) { data: toRunStepDto(ctx.runStep) }); - const toolCall = ctx.toolCall; - if (toolCall) { - if ( - ctx.run.toolApprovals?.find( - (approval) => - approval.toolId === - (toolCall.type === ToolType.USER - ? toolCall.tool.id - : toolCall.type === ToolType.SYSTEM - ? toolCall.toolId - : toolCall.type) - )?.requireApproval === ToolApprovalType.ALWAYS - ) { - await withRedisClient( - (client) => - new Promise((resolve, reject) => { - client.subscribe(createApproveChannel(ctx.run, toolCall), async (err) => { - try { - if (err) { - reject(err); - } else { - ctx.run.requireAction( - new RequiredToolApprove({ - toolCalls: [...(ctx.run.requiredAction?.toolCalls ?? []), toolCall] - }) - ); - await ORM.em.flush(); - await ctx.publish({ - event: 'thread.run.requires_action', - data: toRunDto(ctx.run) - }); - await ctx.publish({ - event: 'done', - data: '[DONE]' - }); - } - } catch (err) { - reject(err); - } - }); - client.on('message', async (_, approval) => { - try { - ctx.run.submitAction(); - await ORM.em.flush(); - if (approval !== 'true') { - reject( - new ToolError('User has not approved this tool to run.', [], { - isFatal: false, - isRetryable: false - }) - ); - } - resolve(true); - } catch (err) { - reject(err); - } - }); - }) - ); - } - } + await requireToolApproval(ctx); + + await requireToolInput(ctx, data); await ctx.publish({ event: 'thread.run.step.in_progress', diff --git a/src/runs/execution/tools/helpers.ts b/src/runs/execution/tools/helpers.ts index 8d89dfc7..6e59744d 100644 --- a/src/runs/execution/tools/helpers.ts +++ b/src/runs/execution/tools/helpers.ts @@ -15,10 +15,14 @@ */ import { + AnyTool, + BaseToolRunOptions, AnyTool as FrameworkTool, StringToolOutput, + ToolError, ToolOutput } from 'bee-agent-framework/tools/base'; +import { setProp } from 'bee-agent-framework/internals/helpers/object'; import { PythonTool } from 'bee-agent-framework/tools/python/python'; import { PythonToolOutput } from 'bee-agent-framework/tools/python/output'; import { Loaded, ref } from '@mikro-orm/core'; @@ -38,7 +42,10 @@ import { LLMTool } from 'bee-agent-framework/tools/llm'; import { AgentContext } from '../execute.js'; import { getRunVectorStores } from '../helpers.js'; -import { CodeInterpreterTool as CodeInterpreterUserTool } from '../../../tools/entities/tool/code-interpreter-tool.entity.js'; +import { + CodeInterpreterTool, + CodeInterpreterTool as CodeInterpreterUserTool +} from '../../../tools/entities/tool/code-interpreter-tool.entity.js'; import { ApiTool as ApiCallUserTool } from '../../../tools/entities/tool/api-tool.entity.js'; import { RedisCache } from '../cache.js'; @@ -76,8 +83,14 @@ import { File } from '@/files/entities/file.entity.js'; import { Attachment } from '@/messages/attachment.entity.js'; import { SystemResource } from '@/tools/entities/tool-resources/system-resource.entity.js'; import { createSearchTool } from '@/runs/execution/tools/search-tool'; -import { sharedRedisCacheClient } from '@/redis.js'; +import { sharedRedisCacheClient, withRedisClient } from '@/redis.js'; import { defaultAIProvider } from '@/runs/execution/provider'; +import { ToolSecret } from '@/tools/entities/tool-secret.entity.js'; +import { createApproveChannel, createToolInputChannel, toRunDto } from '@/runs/runs.service.js'; +import { RequiredToolInput } from '@/runs/entities/requiredToolInput.entity.js'; +import { ToolApprovalType } from '@/runs/entities/toolApproval.entity.js'; +import { RequiredToolApprove } from '@/runs/entities/requiredToolApprove.entity.js'; +import decrypt from '@/utils/crypto/decrypt.js'; const searchCache: SearchToolOptions['cache'] = new RedisCache({ client: sharedRedisCacheClient, @@ -438,3 +451,156 @@ export async function finalizeToolCall( throw new Error(`Unexpected tool call`); } } + +export async function requireToolApproval(ctx: AgentContext) { + const { toolCall } = ctx; + if (toolCall) { + if ( + ctx.run.toolApprovals?.find( + (approval) => + approval.toolId === + (toolCall.type === ToolType.USER + ? toolCall.tool.id + : toolCall.type === ToolType.SYSTEM + ? toolCall.toolId + : toolCall.type) + )?.requireApproval === ToolApprovalType.ALWAYS + ) { + await withRedisClient( + (client) => + new Promise((resolve, reject) => { + client.subscribe(createApproveChannel(ctx.run, toolCall), async (err) => { + try { + if (err) { + reject(err); + } else { + ctx.run.requireAction( + new RequiredToolApprove({ + toolCalls: [...(ctx.run.requiredAction?.toolCalls ?? []), toolCall] + }) + ); + await ORM.em.flush(); + await ctx.publish({ + event: 'thread.run.requires_action', + data: toRunDto(ctx.run) + }); + await ctx.publish({ + event: 'done', + data: '[DONE]' + }); + } + } catch (err) { + reject(err); + } + }); + client.on('message', async (_, approval) => { + try { + ctx.run.submitAction(); + await ORM.em.flush(); + if (approval !== 'true') { + reject( + new ToolError('User has not approved this tool to run.', [], { + isFatal: false, + isRetryable: false + }) + ); + } + resolve(true); + } catch (err) { + reject(err); + } + }); + }) + ); + } + } +} + +export async function requireToolInput( + ctx: AgentContext, + { tool: frameworkTool, options }: { tool: AnyTool; options: BaseToolRunOptions } +) { + const { toolCall } = ctx; + if (toolCall) { + if (toolCall.type === 'user') { + const tool = await ORM.em.getRepository(Tool).findOneOrFail(toolCall.tool.id); + + if (tool instanceof CodeInterpreterTool) { + const toolSecrets = await ORM.em + .getRepository(ToolSecret) + .find({ tool: toolCall.tool.id, createdBy: ctx.run.createdBy, project: ctx.run.project }); + const fulfilledSecrets: { [key: string]: string } = (tool.secrets ?? []).reduce( + (acc, secretName) => { + const secret = toolSecrets.find((ts) => ts.name === secretName); + return { + ...acc, + [secretName]: secret ? decrypt(secret.value) : undefined + }; + }, + {} + ); + const missingSecrets = Object.entries(fulfilledSecrets) + .filter(([_, value]) => !value) + .map(([key]) => key); + if (missingSecrets.length > 0) + await withRedisClient( + (client) => + new Promise((resolve, reject) => { + client.subscribe(createToolInputChannel(ctx.run, toolCall), async (err) => { + try { + if (err) { + reject(err); + } else { + ctx.run.requireAction( + new RequiredToolInput({ + toolCalls: [...(ctx.run.requiredAction?.toolCalls ?? []), toolCall], + inputFields: missingSecrets + }) + ); + await ORM.em.flush(); + await ctx.publish({ + event: 'thread.run.requires_action', + data: toRunDto(ctx.run) + }); + await ctx.publish({ + event: 'done', + data: '[DONE]' + }); + } + } catch (err) { + reject(err); + } + }); + client.on('message', async (_, inputs) => { + try { + ctx.run.submitAction(); + await ORM.em.flush(); + const newSecrest = JSON.parse(inputs); + + newSecrest.forEach((secret: { name: string; value: string }) => { + fulfilledSecrets[secret.name] = secret.value; + }); + + if (frameworkTool instanceof CustomTool) { + Object.entries(fulfilledSecrets).forEach(([key, value]) => + setProp(options, ['env', key], value) + ); + } else { + reject( + new ToolError('Invalid tool type', [], { + isFatal: true, + isRetryable: false + }) + ); + } + resolve(true); + } catch (err) { + reject(err); + } + }); + }) + ); + } + } + } +} diff --git a/src/runs/runs.module.ts b/src/runs/runs.module.ts index 0cfd6c26..0a8e48cc 100644 --- a/src/runs/runs.module.ts +++ b/src/runs/runs.module.ts @@ -33,6 +33,7 @@ import { readRun, readRunTrace, submitToolApproval, + submitToolInputs, submitToolOutput, updateRun } from './runs.service.js'; @@ -79,6 +80,14 @@ import { runSubmitToolApprovalsResponseSchema, runSubmitToolApprovalsStreamSchema } from './dtos/run-submit-tool-approvals.js'; +import { + runSubmitToolInputsBodySchema, + RunSubmitToolInputsParams, + runSubmitToolInputsParamsSchema, + RunSubmitToolInputsBody, + runSubmitToolInputsResponseSchema, + runSubmitToolInputsStreamSchema +} from './dtos/run-submit-tool-inputs.js'; import { Tag } from '@/swagger.js'; @@ -251,7 +260,7 @@ export const runsModule: FastifyPluginAsyncJsonSchemaToTs = async (app) => { } } }, - tags: [Tag.OPENAI_ASSISTANTS_API] + tags: [Tag.BEE_API] }, preHandler: app.auth(), config: { @@ -262,4 +271,30 @@ export const runsModule: FastifyPluginAsyncJsonSchemaToTs = async (app) => { }, async (req) => submitToolApproval({ ...req.params, ...req.body }) ); + + app.post<{ Params: RunSubmitToolInputsParams; Body: RunSubmitToolInputsBody }>( + '/threads/:thread_id/runs/:run_id/submit_tool_inputs', + { + schema: { + params: runSubmitToolInputsParamsSchema, + body: runSubmitToolInputsBodySchema, + response: { + [StatusCodes.OK]: { + content: { + 'application/json': { schema: runSubmitToolInputsResponseSchema }, + 'text/event-stream': { schema: runSubmitToolInputsStreamSchema } + } + } + }, + tags: [Tag.BEE_API] + }, + preHandler: app.auth(), + config: { + rateLimit: { + max: 26 // global rate limit +1 for the UI + } + } + }, + async (req) => submitToolInputs({ ...req.params, ...req.body }) + ); }; diff --git a/src/runs/runs.service.ts b/src/runs/runs.service.ts index c92f680a..ed0f5b16 100644 --- a/src/runs/runs.service.ts +++ b/src/runs/runs.service.ts @@ -16,6 +16,7 @@ import { FilterQuery, Loaded, QueryOrder, ref } from '@mikro-orm/core'; import { Redis } from 'ioredis'; +import { intersection } from 'remeda'; import { RunCreateBody, RunCreateParams, RunCreateResponse } from './dtos/run-create.js'; import { Run, RunStatus } from './entities/run.entity.js'; @@ -42,6 +43,12 @@ import { RequiredToolOutput } from './entities/requiredToolOutput.entity.js'; import { RequiredToolApprove } from './entities/requiredToolApprove.entity.js'; import { ToolApproval, ToolApprovalType } from './entities/toolApproval.entity.js'; import { RequiredActionType } from './entities/requiredAction.entity.js'; +import { RequiredToolInput } from './entities/requiredToolInput.entity.js'; +import { + RunSubmitToolInputsBody, + RunSubmitToolInputsParams, + RunSubmitToolInputsResponse +} from './dtos/run-submit-tool-inputs.js'; import { ORM } from '@/database.js'; import { Thread } from '@/threads/thread.entity.js'; @@ -70,6 +77,7 @@ import { getProjectPrincipal } from '@/administration/helpers.js'; import { RUNS_QUOTA_DAILY } from '@/config.js'; import { dayjs, getLatestDailyFixedTime } from '@/utils/datetime.js'; import { updateRateLimitHeadersWithDailyQuota } from '@/utils/rate-limit.js'; +import { UserCall } from '@/tools/entities/tool-calls/user-call.entity.js'; export async function assertRunsQuota(newRuns = 1) { const count = await ORM.em.getRepository(Run).count({ @@ -108,7 +116,15 @@ export function toRunDto(run: Loaded): RunDto { type: 'submit_tool_approvals', submit_tool_approvals: { tool_calls: run.requiredAction.toolCalls.map(toToolCallDto) } } - : null, + : run.requiredAction instanceof RequiredToolInput + ? { + type: 'submit_tool_inputs', + submit_tool_inputs: { + tool_calls: run.requiredAction.toolCalls.map(toToolCallDto), + input_fields: run.requiredAction.inputFields + } + } + : null, tools: run.tools.map(toToolUsageDto) ?? [], instructions: run.instructions ?? null, additional_instructions: run.additionalInstructions ?? null, @@ -492,6 +508,67 @@ export async function submitToolApproval({ return continueRun({ stream, run, submit }); } +export function createToolInputChannel(run: Run, toolCall: ToolCall) { + return `run:${run.id}:call:${toolCall.id}:input`; +} + +export async function submitToolInputs({ + thread_id, + run_id, + tool_inputs, + stream +}: RunSubmitToolInputsParams & + RunSubmitToolInputsBody): Promise { + const run = await ORM.em.getRepository(Run).findOneOrFail({ id: run_id, thread: thread_id }); + + if ( + run.status !== RunStatus.REQUIRES_ACTION || + !run.requiredAction || + run.requiredAction?.type !== RequiredActionType.INPUT + ) + throw new APIError({ + message: 'No input is required for the run', + code: APIErrorCode.INVALID_INPUT + }); + + const suppliedToolCalls = tool_inputs.map(({ tool_call_id, inputs }) => { + const toolCall = run.requiredAction?.toolCalls.find(({ id }) => id === tool_call_id); + if (!toolCall || !(toolCall instanceof UserCall)) + throw new APIError({ + message: `Unexpected tool call ${tool_call_id}`, + code: APIErrorCode.INVALID_INPUT + }); + + if ( + intersection( + inputs.map((i) => i.name), + (run.requiredAction as RequiredToolInput).inputFields + ).length !== inputs.length + ) { + throw new APIError({ + message: `Missing required tool input. Provide all input fields: ${(run.requiredAction as RequiredToolInput).inputFields}`, + code: APIErrorCode.INVALID_INPUT + }); + } + return { toolCall, inputs }; + }); + if (suppliedToolCalls.length < run.requiredAction.toolCalls.length) + throw new APIError({ + message: 'Missing tool calls', + code: APIErrorCode.INVALID_INPUT + }); + + const submit = async (client: Redis) => { + await Promise.all( + suppliedToolCalls.map(({ toolCall, inputs }) => { + client.publish(createToolInputChannel(run, toolCall), JSON.stringify(inputs)); + }) + ); + }; + + return continueRun({ stream, run, submit }); +} + async function continueRun({ stream, run, diff --git a/src/server.ts b/src/server.ts index 36de2d96..82c73e04 100644 --- a/src/server.ts +++ b/src/server.ts @@ -49,6 +49,7 @@ import { projectsModule } from './administration/projects.module.js'; import { projectUsersModule } from './administration/project-users.module.js'; import { organizationUsersModule } from './administration/organization-users.module.js'; import { apiKeysModule } from './administration/api-keys.module.js'; +import { toolSecretsModule } from './tools/tool-secrets.module.js'; import { artifactsModule } from './artifacts/artifacts.module.js'; import { chatModule } from './chat/chat.module.js'; import { embeddingsModule } from './embeddings/embeddings.module.js'; @@ -88,6 +89,7 @@ try { app.register(messagesModule, { prefix: '/v1' }); app.register(runsModule, { prefix: '/v1' }); app.register(toolsModule, { prefix: '/v1' }); + app.register(toolSecretsModule, { prefix: '/v1' }); app.register(runStepsModule, { prefix: '/v1' }); app.register(filesModule, { prefix: '/v1' }); app.register(vectorStoresModule, { prefix: '/v1' }); diff --git a/src/tools/dtos/tool-create.ts b/src/tools/dtos/tool-create.ts index e369cafd..eb2a2732 100644 --- a/src/tools/dtos/tool-create.ts +++ b/src/tools/dtos/tool-create.ts @@ -30,7 +30,8 @@ export const toolCreateBodySchema = { name: { type: 'string' }, source_code: { type: 'string' }, metadata: metadataSchema, - user_description: { type: 'string' } + user_description: { type: 'string' }, + secrets: { type: 'array', items: { type: 'string' } } } }, { diff --git a/src/tools/dtos/tool-secret-create.ts b/src/tools/dtos/tool-secret-create.ts new file mode 100644 index 00000000..cf94d16e --- /dev/null +++ b/src/tools/dtos/tool-secret-create.ts @@ -0,0 +1,34 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { FromSchema, JSONSchema } from 'json-schema-to-ts'; + +import { toolSecretSchema } from './tool-secret'; + +export const toolSecretCreateBodySchema = { + type: 'object', + additionalProperties: false, + required: ['name', 'value', 'tool_id'], + properties: { + name: { type: 'string' }, + value: { type: 'string' }, + tool_id: { type: 'string' } + } +} as const satisfies JSONSchema; +export type ToolSecretCreateBody = FromSchema; + +export const toolSecretCreateResponseSchema = toolSecretSchema; +export type ToolSecretCreateResponse = FromSchema; diff --git a/src/tools/dtos/tool-secret-delete.ts b/src/tools/dtos/tool-secret-delete.ts new file mode 100644 index 00000000..7b9ee072 --- /dev/null +++ b/src/tools/dtos/tool-secret-delete.ts @@ -0,0 +1,27 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { FromSchema } from 'json-schema-to-ts'; + +import { toolSecretReadParamsSchema } from './tool-secret-read'; + +import { createDeleteSchema } from '@/schema.js'; + +export const toolSecretDeleteParamsSchema = toolSecretReadParamsSchema; +export type ToolSecretDeleteParams = FromSchema; + +export const toolSecretDeleteResponseSchema = createDeleteSchema('tool-secret'); +export type ToolSecretDeleteResponse = FromSchema; diff --git a/src/tools/dtos/tool-secret-read.ts b/src/tools/dtos/tool-secret-read.ts new file mode 100644 index 00000000..f53fcbaf --- /dev/null +++ b/src/tools/dtos/tool-secret-read.ts @@ -0,0 +1,34 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { FromSchema, JSONSchema } from 'json-schema-to-ts'; + +import { toolSecretSchema } from './tool-secret'; + +export const toolSecretReadParamsSchema = { + type: 'object', + additionalProperties: false, + required: ['tool_secret_id'], + properties: { + tool_secret_id: { + type: 'string' + } + } +} as const satisfies JSONSchema; +export type ToolSecretReadParams = FromSchema; + +export const toolSecretReadResponseSchema = toolSecretSchema; +export type ToolSecretReadResponse = FromSchema; diff --git a/src/tools/dtos/tool-secret-update.ts b/src/tools/dtos/tool-secret-update.ts new file mode 100644 index 00000000..c1eda70a --- /dev/null +++ b/src/tools/dtos/tool-secret-update.ts @@ -0,0 +1,38 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { FromSchema, JSONSchema } from 'json-schema-to-ts'; + +import { toolSecretReadParamsSchema } from './tool-secret-read'; +import { toolSecretSchema } from './tool-secret'; + +export const toolSecretUpdateBodySchema = { + type: 'object', + additionalProperties: false, + required: ['name', 'value', 'tool_id'], + properties: { + name: { type: 'string' }, + value: { type: 'string' }, + tool_id: { type: 'string' } + } +} as const satisfies JSONSchema; +export type ToolSecretUpdateBody = FromSchema; + +export const toolSecretUpdateParamsSchema = toolSecretReadParamsSchema; +export type ToolSecretUpdateParams = FromSchema; + +export const toolSecretUpdateResponseSchema = toolSecretSchema; +export type ToolSecretUpdateResponse = FromSchema; diff --git a/src/tools/dtos/tool-secret.ts b/src/tools/dtos/tool-secret.ts new file mode 100644 index 00000000..348f67ca --- /dev/null +++ b/src/tools/dtos/tool-secret.ts @@ -0,0 +1,38 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { FromSchema, JSONSchema } from 'json-schema-to-ts'; + +export const toolSecretSchema = { + type: 'object', + required: ['id', 'object', 'name', 'value', 'tool_id'], + properties: { + id: { + type: 'string' + }, + object: { const: 'tool-secret' }, + name: { + type: 'string' + }, + value: { + type: 'string' + }, + tool_id: { + type: 'string' + } + } +} as const satisfies JSONSchema; +export type ToolSecret = FromSchema; diff --git a/src/tools/dtos/tool-secrets-list.ts b/src/tools/dtos/tool-secrets-list.ts new file mode 100644 index 00000000..2fe9b115 --- /dev/null +++ b/src/tools/dtos/tool-secrets-list.ts @@ -0,0 +1,27 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { FromSchema } from 'json-schema-to-ts'; + +import { toolSecretSchema } from './tool-secret'; + +import { createPaginationQuerySchema, withPagination } from '@/schema.js'; + +export const toolSecretsListQuerySchema = createPaginationQuerySchema(); +export type ToolSecretsListQuery = FromSchema; + +export const toolSecretsListResponseSchema = withPagination(toolSecretSchema); +export type ToolSecretsListResponse = FromSchema; diff --git a/src/tools/dtos/tool-update.ts b/src/tools/dtos/tool-update.ts index ad31d083..1441d06f 100644 --- a/src/tools/dtos/tool-update.ts +++ b/src/tools/dtos/tool-update.ts @@ -36,6 +36,7 @@ export const toolUpdateBodySchema = { source_code: { type: 'string' }, + secrets: { type: 'array', items: { type: 'string' } }, metadata: metadataSchema, user_description: { type: 'string' } } diff --git a/src/tools/entities/tool-secret.entity.ts b/src/tools/entities/tool-secret.entity.ts new file mode 100644 index 00000000..76cc3989 --- /dev/null +++ b/src/tools/entities/tool-secret.entity.ts @@ -0,0 +1,58 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Entity, Index, ManyToOne, Property, Ref } from '@mikro-orm/core'; + +import { Tool } from './tool/tool.entity'; + +import { ProjectScopedEntity, ProjectScopedEntityInput } from '@/common/project-scoped.entity'; + +@Entity() +@Index({ + options: [ + { + name: 1, + tool: 1 + }, + { + unique: true, + partialFilterExpression: { deletedAt: { $in: [null] } } + } + ] +}) +export class ToolSecret extends ProjectScopedEntity { + getIdPrefix(): string { + return 'tool-secret'; + } + + @Property() + name: string; + + @Property() + value: string; + + @ManyToOne() + tool: Ref; + + constructor({ name, value, tool, ...rest }: ToolSecretInput) { + super(rest); + this.name = name; + this.value = value; + this.tool = tool; + } +} + +type ToolSecretInput = ProjectScopedEntityInput & Pick; diff --git a/src/tools/entities/tool/code-interpreter-tool.entity.ts b/src/tools/entities/tool/code-interpreter-tool.entity.ts index 8ded2a3d..8740c300 100644 --- a/src/tools/entities/tool/code-interpreter-tool.entity.ts +++ b/src/tools/entities/tool/code-interpreter-tool.entity.ts @@ -29,12 +29,17 @@ export class CodeInterpreterTool extends Tool { @Property({ type: 'json' }) jsonSchema!: JSONSchema; - constructor({ sourceCode, jsonSchema, ...rest }: CodeInterpreterToolInput) { + @Property() + secrets?: string[]; + + constructor({ sourceCode, jsonSchema, secrets, ...rest }: CodeInterpreterToolInput) { super(rest); this.sourceCode = sourceCode; this.jsonSchema = jsonSchema; + this.secrets = secrets; } } export type CodeInterpreterToolInput = ToolInput & - Pick; + Pick & + Partial>; diff --git a/src/tools/tool-secrets.module.ts b/src/tools/tool-secrets.module.ts new file mode 100644 index 00000000..caf1b53b --- /dev/null +++ b/src/tools/tool-secrets.module.ts @@ -0,0 +1,133 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { FastifyPluginAsyncJsonSchemaToTs } from '@fastify/type-provider-json-schema-to-ts'; +import { StatusCodes } from 'http-status-codes'; + +import { + ToolSecretCreateBody, + toolSecretCreateBodySchema, + toolSecretCreateResponseSchema +} from './dtos/tool-secret-create'; +import { + ToolSecretUpdateBody, + toolSecretUpdateBodySchema, + ToolSecretUpdateParams, + toolSecretUpdateParamsSchema, + toolSecretUpdateResponseSchema +} from './dtos/tool-secret-update'; +import { + ToolSecretsListQuery, + toolSecretsListQuerySchema, + toolSecretsListResponseSchema +} from './dtos/tool-secrets-list'; +import { + ToolSecretReadParams, + toolSecretReadParamsSchema, + toolSecretReadResponseSchema +} from './dtos/tool-secret-read'; +import { + ToolSecretDeleteParams, + toolSecretDeleteParamsSchema, + toolSecretDeleteResponseSchema +} from './dtos/tool-secret-delete'; +import { + createToolSecret, + deleteToolSecret, + listToolSecrets, + readToolSecret, + updateToolSecret +} from './tool-secrets.service'; + +import { Tag } from '@/swagger.js'; + +export const toolSecretsModule: FastifyPluginAsyncJsonSchemaToTs = async (app) => { + app.post<{ Body: ToolSecretCreateBody }>( + '/tool_secrets', + { + preHandler: app.auth(), + schema: { + body: toolSecretCreateBodySchema, + response: { + [StatusCodes.OK]: toolSecretCreateResponseSchema + }, + tags: [Tag.BEE_API] + } + }, + async (req) => createToolSecret(req.body) + ); + + app.post<{ Body: ToolSecretUpdateBody; Params: ToolSecretUpdateParams }>( + '/tool_secrets/:tool_secret_id', + { + preHandler: app.auth(), + schema: { + body: toolSecretUpdateBodySchema, + params: toolSecretUpdateParamsSchema, + response: { + [StatusCodes.OK]: toolSecretUpdateResponseSchema + }, + tags: [Tag.BEE_API] + } + }, + async (req) => updateToolSecret({ ...req.body, ...req.params }) + ); + + app.get<{ Querystring: ToolSecretsListQuery }>( + '/tool_secrets', + { + preHandler: app.auth(), + schema: { + querystring: toolSecretsListQuerySchema, + response: { + [StatusCodes.OK]: toolSecretsListResponseSchema + }, + tags: [Tag.BEE_API] + } + }, + async (req) => listToolSecrets(req.query) + ); + + app.get<{ Params: ToolSecretReadParams }>( + '/tool_secrets/:tool_secret_id', + { + preHandler: app.auth(), + schema: { + params: toolSecretReadParamsSchema, + response: { + [StatusCodes.OK]: toolSecretReadResponseSchema + }, + tags: [Tag.BEE_API] + } + }, + async (req) => readToolSecret(req.params) + ); + + app.delete<{ Params: ToolSecretDeleteParams }>( + '/tool_secrets/:tool_secret_id', + { + preHandler: app.auth(), + schema: { + params: toolSecretDeleteParamsSchema, + response: { + [StatusCodes.OK]: toolSecretDeleteResponseSchema + }, + tags: [Tag.BEE_API] + } + }, + async (req) => deleteToolSecret(req.params) + ); +}; diff --git a/src/tools/tool-secrets.service.ts b/src/tools/tool-secrets.service.ts new file mode 100644 index 00000000..09bd902b --- /dev/null +++ b/src/tools/tool-secrets.service.ts @@ -0,0 +1,117 @@ +/** + * Copyright 2024 IBM Corp. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Loaded, ref } from '@mikro-orm/core'; + +import { Tool } from './entities/tool/tool.entity.js'; +import { ToolSecret } from './entities/tool-secret.entity.js'; +import { ToolSecret as ToolSecretDto } from './dtos/tool-secret.js'; +import { ToolSecretsListQuery, ToolSecretsListResponse } from './dtos/tool-secrets-list.js'; +import { ToolSecretCreateBody, ToolSecretCreateResponse } from './dtos/tool-secret-create.js'; +import { + ToolSecretUpdateBody, + ToolSecretUpdateParams, + ToolSecretUpdateResponse +} from './dtos/tool-secret-update.js'; +import { ToolSecretReadParams, ToolSecretReadResponse } from './dtos/tool-secret-read.js'; +import { ToolSecretDeleteParams, ToolSecretDeleteResponse } from './dtos/tool-secret-delete.js'; + +import { createDeleteResponse } from '@/utils/delete.js'; +import { ORM } from '@/database.js'; +import { createPaginatedResponse, getListCursor } from '@/utils/pagination.js'; +import { getUpdatedValue } from '@/utils/update.js'; +import { redactKey } from '@/administration/helpers.js'; +import encrypt from '@/utils/crypto/encrypt.js'; +import decrypt from '@/utils/crypto/decrypt.js'; + +export function toToolSecretDto(toolSecret: Loaded): ToolSecretDto { + return { + id: toolSecret.id, + object: 'tool-secret', + value: redactKey(decrypt(toolSecret.value)), + name: toolSecret.name, + tool_id: toolSecret.tool.id + }; +} + +export async function listToolSecrets({ + limit, + after, + before, + order, + order_by +}: ToolSecretsListQuery): Promise { + const cursor = await getListCursor>( + {}, + { limit, order, order_by, after, before }, + ORM.em.getRepository(ToolSecret) + ); + return createPaginatedResponse(cursor, toToolSecretDto); +} + +export async function createToolSecret({ + name, + value, + tool_id +}: ToolSecretCreateBody): Promise { + const tool = await ORM.em.getRepository(Tool).findOneOrFail({ id: tool_id }); + + const toolSecret = new ToolSecret({ + name, + value: encrypt(value), + tool: ref(tool) + }); + + await ORM.em.persistAndFlush(toolSecret); + + return toToolSecretDto(toolSecret); +} + +export async function updateToolSecret({ + tool_id, + ...body +}: ToolSecretUpdateBody & ToolSecretUpdateParams): Promise { + const toolSecret = await ORM.em.getRepository(ToolSecret).findOneOrFail({ id: tool_id }); + + toolSecret.name = getUpdatedValue(body.name, toolSecret.name); + toolSecret.value = getUpdatedValue(body.value && encrypt(body.value), toolSecret.value); + if (tool_id) { + const tool = await ORM.em.getRepository(Tool).findOneOrFail({ id: tool_id }); + toolSecret.tool = ref(tool); + } + + await ORM.em.flush(); + + return toToolSecretDto(toolSecret); +} + +export async function readToolSecret({ + tool_secret_id +}: ToolSecretReadParams): Promise { + const toolSecret = await ORM.em.getRepository(ToolSecret).findOneOrFail({ id: tool_secret_id }); + return toToolSecretDto(toolSecret); +} + +export async function deleteToolSecret({ + tool_secret_id +}: ToolSecretDeleteParams): Promise { + const toolSecret = await ORM.em.getRepository(ToolSecret).findOneOrFail({ id: tool_secret_id }); + + toolSecret.delete(); + await ORM.em.flush(); + + return createDeleteResponse(tool_secret_id, 'tool-secret'); +} diff --git a/src/tools/tools.service.ts b/src/tools/tools.service.ts index 4997d896..0726637c 100644 --- a/src/tools/tools.service.ts +++ b/src/tools/tools.service.ts @@ -691,6 +691,7 @@ async function createCodeInterpreterTool( sourceCode: body.source_code, jsonSchema: customTool.inputSchema(), description: customTool.description, + secrets: body.secrets, metadata: body.metadata ?? undefined, userDescription: body.user_description }); @@ -833,6 +834,7 @@ export async function updateTool({ ); tool.jsonSchema = newCustomTool.inputSchema(); tool.description = newCustomTool.description; + tool.secrets = getUpdatedValue(body.secrets, tool.secrets); } catch (err) { if (err instanceof CustomToolCreateError) { throw new APIError({