diff --git a/package.json b/package.json index 160a609..eb6bc5f 100644 --- a/package.json +++ b/package.json @@ -48,8 +48,8 @@ "@zilliz/milvus2-sdk-node": "^2.4.4", "ajv": "^8.17.1", "axios": "^1.7.7", - "bee-agent-framework": "0.0.42", - "bee-observe-connector": "0.0.5", + "bee-agent-framework": "0.0.44", + "bee-observe-connector": "0.0.6", "bullmq": "5.8.1", "cache-manager": "^5.7.6", "dayjs": "^1.11.11", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index d41d2bc..765a043 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -84,11 +84,11 @@ importers: specifier: ^1.7.7 version: 1.7.7 bee-agent-framework: - specifier: 0.0.42 - version: 0.0.42(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8)) + specifier: 0.0.44 + version: 0.0.44(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(@zilliz/milvus2-sdk-node@2.4.4)(google-auth-library@9.15.0)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8)) bee-observe-connector: - specifier: 0.0.5 - version: 0.0.5(bee-agent-framework@0.0.42(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8))) + specifier: 0.0.6 + version: 0.0.6(bee-agent-framework@0.0.44(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(@zilliz/milvus2-sdk-node@2.4.4)(google-auth-library@9.15.0)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8))) bullmq: specifier: 5.8.1 version: 5.8.1 @@ -1971,17 +1971,20 @@ packages: resolution: {integrity: sha512-4Bcg1P8xhUuqcii/S0Z9wiHIrQVPMermM1any+MX5GeGD7faD3/msQUDGLol9wOcz4/jbg/WJnGqoJF6LiBdtg==} engines: {node: '>=10.0.0'} - bee-agent-framework@0.0.42: - resolution: {integrity: sha512-W6onNf9Zaj6GNpRTKJOMChfvSZccrBsmafuSh8I7Er1eSND9viu5RVybG/FiAI6omjdJGJnEZ+lTgY46rl+DdA==} + bee-agent-framework@0.0.44: + resolution: {integrity: sha512-62K8lqrHxzD8rYVGx+MqY2y42B19MpMZ30Yy0x4G4eNMOuyA0uhp2EsAlPwENxFRjP6Yn1bPXVntaWhCDvY5+g==} peerDependencies: '@aws-sdk/client-bedrock-runtime': ^3.687.0 '@elastic/elasticsearch': ^8.0.0 + '@google-cloud/vertexai': '*' '@googleapis/customsearch': ^3.2.0 '@grpc/grpc-js': ^1.11.3 '@grpc/proto-loader': ^0.7.13 '@ibm-generative-ai/node-sdk': ~3.2.4 '@langchain/community': '>=0.2.28' '@langchain/core': '>=0.2.27' + '@zilliz/milvus2-sdk-node': ^2.4.9 + google-auth-library: '*' groq-sdk: ^0.7.0 ollama: ^0.5.8 openai: ^4.67.3 @@ -1992,6 +1995,8 @@ packages: optional: true '@elastic/elasticsearch': optional: true + '@google-cloud/vertexai': + optional: true '@googleapis/customsearch': optional: true '@grpc/grpc-js': @@ -2004,6 +2009,10 @@ packages: optional: true '@langchain/core': optional: true + '@zilliz/milvus2-sdk-node': + optional: true + google-auth-library: + optional: true groq-sdk: optional: true ollama: @@ -2015,8 +2024,8 @@ packages: sequelize: optional: true - bee-observe-connector@0.0.5: - resolution: {integrity: sha512-veiHn7u4sEPUgMMrNCJ/yuS6hVPE/qSxIJzaGAyolEsF0HXIxP8HdZdGkkZPrq0fUVXXiKBCwIQtt7cAnpYT8g==} + bee-observe-connector@0.0.6: + resolution: {integrity: sha512-KZhhOFwvw/dR2YN4DPf7dZkquJ4O41D0WhjAsU3+RApkvBnpSm6Y/TYTa7vOZKaJzbuJ0NqRVrhPD8dCuKXFMQ==} peerDependencies: bee-agent-framework: '>=0.0.27 <0.1.0' @@ -2513,8 +2522,8 @@ packages: resolution: {integrity: sha512-ZmdL2rui+eB2YwhsWzjInR8LldtZHGDoQ1ugH85ppHKwpUHL7j7rN0Ti9NCnGiQbhaZ11FpR+7ao1dNsmduNUg==} engines: {node: '>=12'} - duck-duck-scrape@2.2.5: - resolution: {integrity: sha512-RTu/Ag5LhgD/j1l2zVGCwTINBoEYCffl58nMoBjtXWJG8tTex72h3gxpDjAr8jVqFaBvhgASNKTUsE31JeqYgw==} + duck-duck-scrape@2.2.6: + resolution: {integrity: sha512-ov8IwoueOXnXsqGrN6ESys9x2uesyCVZYRex9is9uXyDuI3Je9pxjH/Y6atIUVjBhnwWQ5d1omoo2vbUmVDBLQ==} eastasianwidth@0.2.0: resolution: {integrity: sha512-I88TYZWc9XiYHRQ4/3c5rjjfgkjhLyW2luGIheGERbNQ6OY7yTybanSpDXZa8y7VUP9YmDcYa+eyq4ca7iLqWA==} @@ -7260,7 +7269,7 @@ snapshots: basic-ftp@5.0.5: {} - bee-agent-framework@0.0.42(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8)): + bee-agent-framework@0.0.44(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(@zilliz/milvus2-sdk-node@2.4.4)(google-auth-library@9.15.0)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8)): dependencies: '@ai-zen/node-fetch-event-source': 2.1.4 '@connectrpc/connect': 1.6.1(@bufbuild/protobuf@1.10.0) @@ -7270,7 +7279,7 @@ snapshots: ajv: 8.17.1 ajv-formats: 3.0.1(ajv@8.17.1) bee-proto: 0.0.2 - duck-duck-scrape: 2.2.5 + duck-duck-scrape: 2.2.6 fast-xml-parser: 4.5.0 header-generator: 2.1.56 joplin-turndown-plugin-gfm: 1.0.12 @@ -7297,6 +7306,8 @@ snapshots: '@grpc/grpc-js': 1.12.2 '@grpc/proto-loader': 0.7.13 '@ibm-generative-ai/node-sdk': 3.2.4 + '@zilliz/milvus2-sdk-node': 2.4.4 + google-auth-library: 9.15.0 ollama: 0.5.9 openai: 4.67.3(zod@3.23.8) openai-chat-tokens: 0.2.8 @@ -7305,9 +7316,9 @@ snapshots: - debug - encoding - bee-observe-connector@0.0.5(bee-agent-framework@0.0.42(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8))): + bee-observe-connector@0.0.6(bee-agent-framework@0.0.44(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(@zilliz/milvus2-sdk-node@2.4.4)(google-auth-library@9.15.0)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8))): dependencies: - bee-agent-framework: 0.0.42(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8)) + bee-agent-framework: 0.0.44(@bufbuild/protobuf@1.10.0)(@googleapis/customsearch@3.2.0)(@grpc/grpc-js@1.12.2)(@grpc/proto-loader@0.7.13)(@ibm-generative-ai/node-sdk@3.2.4)(@zilliz/milvus2-sdk-node@2.4.4)(google-auth-library@9.15.0)(ollama@0.5.9)(openai-chat-tokens@0.2.8)(openai@4.67.3(zod@3.23.8)) openapi-fetch: 0.11.3 remeda: 2.16.0 @@ -7816,7 +7827,7 @@ snapshots: dotenv@16.4.5: {} - duck-duck-scrape@2.2.5: + duck-duck-scrape@2.2.6: dependencies: html-entities: 2.5.2 needle: 3.3.1 diff --git a/src/files/extraction/helpers.ts b/src/files/extraction/helpers.ts index aea1d00..1f69af2 100644 --- a/src/files/extraction/helpers.ts +++ b/src/files/extraction/helpers.ts @@ -17,6 +17,7 @@ import { Loaded } from '@mikro-orm/core'; import mime from 'mime'; import { recursiveSplitString } from 'bee-agent-framework/internals/helpers/string'; +import ibm from 'ibm-cos-sdk'; import { s3Client } from '../files.service'; import { DoclingExtraction } from '../entities/extractions/docling-extraction.entity'; @@ -34,6 +35,12 @@ import { EXTRACTION_BACKEND, S3_BUCKET_FILE_STORAGE } from '@/config'; import { ORM } from '@/database'; import { QueueName } from '@/jobs/constants'; +export const withAbort = (value: ibm.Request, signal?: AbortSignal) => { + const handler = () => value.abort(); + signal?.addEventListener('abort', handler); + return value.promise().finally(() => signal?.removeEventListener('abort', handler)); +}; + function isNativeDoclingFormat(mimeType: string): boolean { const extension = mime.getExtension(mimeType); if (!extension) return false; @@ -200,81 +207,72 @@ export async function scheduleExtraction( } } -export async function removeExtraction(file: Loaded) { +type AvailableKeys = Exclude; + +const keyByProvider = { + [ExtractionBackend.DOCLING]: ['documentStorageId', 'chunksStorageId', 'textStorageId'], + [ExtractionBackend.WDU]: ['storageId'], + [ExtractionBackend.UNSTRUCTURED_OPENSOURCE]: ['storageId'], + [ExtractionBackend.UNSTRUCTURED_API]: ['storageId'] +} as const satisfies Record[]>; + +export async function removeExtraction(file: Loaded, signal?: AbortSignal) { const extraction = file.extraction; if (!extraction) throw new Error('No extraction to remove'); - switch (extraction.backend) { - case ExtractionBackend.DOCLING: - if (extraction.documentStorageId) - await s3Client - .deleteObject({ Bucket: S3_BUCKET_FILE_STORAGE, Key: extraction.documentStorageId }) - .promise(); - if (extraction.chunksStorageId) - await s3Client - .deleteObject({ Bucket: S3_BUCKET_FILE_STORAGE, Key: extraction.chunksStorageId }) - .promise(); - if (extraction.textStorageId) - await s3Client - .deleteObject({ Bucket: S3_BUCKET_FILE_STORAGE, Key: extraction.textStorageId }) - .promise(); - break; - case ExtractionBackend.WDU: - if (extraction.storageId) - await s3Client - .deleteObject({ Bucket: S3_BUCKET_FILE_STORAGE, Key: extraction.storageId }) - .promise(); - break; - case ExtractionBackend.UNSTRUCTURED_OPENSOURCE: - case ExtractionBackend.UNSTRUCTURED_API: - if (extraction.storageId) - await s3Client - .deleteObject({ Bucket: S3_BUCKET_FILE_STORAGE, Key: extraction.storageId }) - .promise(); - break; - } + + await Promise.all( + keyByProvider[extraction.backend].map(async (property) => { + const Key = extraction[property as keyof typeof extraction]; + if (Key) { + await withAbort(s3Client.deleteObject({ Bucket: S3_BUCKET_FILE_STORAGE, Key }), signal); + } + }) + ); + file.extraction = undefined; await ORM.em.flush(); } -export async function getExtractedText(file: Loaded) { +export async function getExtractedText(file: Loaded, signal?: AbortSignal): Promise { const extraction = file.extraction; if (!extraction) throw new Error('Extraction not found'); switch (extraction.backend) { case ExtractionBackend.WDU: { if (!extraction.storageId) throw new Error('Extraction missing'); - const object = await s3Client - .getObject({ + const object = await withAbort( + s3Client.getObject({ Bucket: S3_BUCKET_FILE_STORAGE, Key: extraction.storageId - }) - .promise(); + }), + signal + ); const body = object.Body; if (!body) throw new Error('Invalid Body of a file'); return body.toString('utf-8'); } case ExtractionBackend.DOCLING: { if (!extraction.textStorageId) throw new Error('Extraction missing'); - return readTextFile(extraction.textStorageId); + return readTextFile(extraction.textStorageId, signal); } case ExtractionBackend.UNSTRUCTURED_OPENSOURCE: case ExtractionBackend.UNSTRUCTURED_API: { if (!extraction.storageId) throw new Error('Extraction missing'); const elements = JSON.parse( - await readTextFile(extraction.storageId) + await readTextFile(extraction.storageId, signal) ) as UnstructuredExtractionDocument; return elements.map((element) => element.text).join(''); } } } -export async function getExtractedChunks(file: Loaded) { +export async function getExtractedChunks(file: Loaded, signal?: AbortSignal) { const extraction = file.extraction; if (!extraction) throw new Error('Extraction not found'); switch (extraction.backend) { case ExtractionBackend.DOCLING: { if (!extraction.chunksStorageId) { if (!extraction.textStorageId) throw new Error('Extraction missing'); - const text = await getExtractedText(file); + const text = await getExtractedText(file, signal); const splitter = recursiveSplitString(text, { size: 400, overlap: 200, @@ -283,12 +281,12 @@ export async function getExtractedChunks(file: Loaded) { return Array.from(splitter); } const chunks = JSON.parse( - await readTextFile(extraction.chunksStorageId) + await readTextFile(extraction.chunksStorageId, signal) ) as DoclingChunksExtraction; return chunks.map((c) => c.text); } case ExtractionBackend.WDU: { - const text = await getExtractedText(file); + const text = await getExtractedText(file, signal); const splitter = recursiveSplitString(text, { size: 400, overlap: 200, @@ -300,7 +298,7 @@ export async function getExtractedChunks(file: Loaded) { case ExtractionBackend.UNSTRUCTURED_API: { if (!extraction.storageId) throw new Error('Extraction missing'); const elements = JSON.parse( - await readTextFile(extraction.storageId) + await readTextFile(extraction.storageId, signal) ) as UnstructuredExtractionDocument; return elements .filter((element) => element.type === 'CompositeElement') @@ -309,13 +307,14 @@ export async function getExtractedChunks(file: Loaded) { } } -async function readTextFile(key: string) { - const object = await s3Client - .getObject({ +async function readTextFile(key: string, signal?: AbortSignal) { + const object = await withAbort( + s3Client.getObject({ Bucket: S3_BUCKET_FILE_STORAGE, Key: key - }) - .promise(); + }), + signal + ); const body = object.Body; if (!body) throw new Error('Invalid Body of a file'); const data = body.toString('utf-8'); diff --git a/src/runs/execution/factory.ts b/src/runs/execution/factory.ts index 2a1d270..87427f2 100644 --- a/src/runs/execution/factory.ts +++ b/src/runs/execution/factory.ts @@ -36,14 +36,13 @@ import { WatsonXLLM } from 'bee-agent-framework/adapters/watsonx/llm'; import { ZodType } from 'zod'; import { PromptTemplate } from 'bee-agent-framework'; import { AnyTool } from 'bee-agent-framework/tools/base'; -import { GraniteBeeAgent } from 'bee-agent-framework/agents/granite/agent'; import { StreamlitAgent } from 'bee-agent-framework/agents/experimental/streamlit/agent'; -import { GraniteBeeSystemPrompt } from 'bee-agent-framework/agents/granite/prompts'; import { BeeAgent } from 'bee-agent-framework/agents/bee/agent'; import { BeeSystemPrompt } from 'bee-agent-framework/agents/bee/prompts'; import { ChatLLM, ChatLLMOutput } from 'bee-agent-framework/llms/chat'; import { BaseMemory } from 'bee-agent-framework/memory/base'; import { StreamlitAgentSystemPrompt } from 'bee-agent-framework/agents/experimental/streamlit/prompts'; +import { GraniteBeeSystemPrompt } from 'bee-agent-framework/agents/bee/runners/granite/prompts'; import { Run } from '../entities/run.entity'; @@ -252,19 +251,17 @@ export function createAgentRun( ] as const; switch (run.assistant.$.agent) { case Agent.BEE: { - const agent = run.model.includes('granite') - ? new GraniteBeeAgent({ - llm, - memory, - tools, - templates: { system: getPromptTemplate(run, GraniteBeeSystemPrompt) } - }) - : new BeeAgent({ - llm, - memory, - tools, - templates: { system: getPromptTemplate(run, BeeSystemPrompt) } - }); + const agent = new BeeAgent({ + llm, + memory, + tools, + templates: { + system: getPromptTemplate( + run, + run.model.includes('granite') ? GraniteBeeSystemPrompt : BeeSystemPrompt + ) + } + }); return [agent.run(...runArgs).observe(createBeeStreamingHandler(ctx)), agent]; } case Agent.STREAMLIT: { diff --git a/src/runs/execution/tools/api-call-tool.ts b/src/runs/execution/tools/api-call-tool.ts index 14c46c2..966301a 100644 --- a/src/runs/execution/tools/api-call-tool.ts +++ b/src/runs/execution/tools/api-call-tool.ts @@ -16,13 +16,20 @@ import { join } from 'path'; -import { BaseToolOptions, StringToolOutput, Tool, ToolError } from 'bee-agent-framework/tools/base'; +import { + BaseToolOptions, + BaseToolRunOptions, + StringToolOutput, + Tool, + ToolError +} from 'bee-agent-framework/tools/base'; import { SchemaObject } from 'ajv'; import { parse } from 'yaml'; import { isEmpty } from 'remeda'; import axios from 'axios'; import { HttpProxyAgent } from 'http-proxy-agent'; import { HttpsProxyAgent } from 'https-proxy-agent'; +import { GetRunContext } from 'bee-agent-framework/context'; import { AgentContext } from '../execute.js'; @@ -103,7 +110,11 @@ export class ApiCallTool extends Tool { } } - protected async _run(input: any) { + protected async _run( + input: any, + _options: BaseToolRunOptions | undefined, + run: GetRunContext + ) { let path: string = input.path || ''; const url = new URL(this.openApiSchema.servers[0].url); Object.keys(input.parameters ?? {}).forEach((key) => { @@ -127,11 +138,11 @@ export class ApiCallTool extends Tool { transformResponse: [(data) => data], httpsAgent: HTTP_PROXY_URL && new HttpsProxyAgent(HTTP_PROXY_URL), httpAgent: HTTP_PROXY_URL && new HttpProxyAgent(HTTP_PROXY_URL), - signal: AbortSignal.timeout(30_000) + signal: AbortSignal.any([AbortSignal.timeout(30_000), run.signal]) }); return new StringToolOutput(response.data); } catch (error) { - throw new ToolError(`Request to ${url} failed.`); + throw new ToolError(`Request to ${url} failed.`, [error]); } } } diff --git a/src/runs/execution/tools/file-search-tool.ts b/src/runs/execution/tools/file-search-tool.ts index 796c9c8..e011518 100644 --- a/src/runs/execution/tools/file-search-tool.ts +++ b/src/runs/execution/tools/file-search-tool.ts @@ -24,6 +24,7 @@ import { import { z } from 'zod'; import { isTruthy } from 'remeda'; import { Loaded } from '@mikro-orm/core'; +import { GetRunContext } from 'bee-agent-framework/context'; import { getVectorStoreClient } from '@/vector-store-files/execution/client.js'; import { VectorStore } from '@/vector-stores/entities/vector-store.entity.js'; @@ -81,13 +82,14 @@ export class FileSearchTool extends Tool, - options: BaseToolRunOptions + _options: BaseToolRunOptions | undefined, + run: GetRunContext ): Promise { const vectorStoreClient = getVectorStoreClient(); const embeddingAdapter = await createEmbeddingAdapter(); - const embedding = await embeddingAdapter.embed(query, { signal: options.signal }); + const embedding = await embeddingAdapter.embed(query, { signal: run.signal }); if (!embedding) throw new Error('Missing embedding data in embedding response'); if (this.vectorStores.some((vectorStore) => vectorStore.expired)) { diff --git a/src/runs/execution/tools/function.ts b/src/runs/execution/tools/function.ts index bade934..db50604 100644 --- a/src/runs/execution/tools/function.ts +++ b/src/runs/execution/tools/function.ts @@ -17,12 +17,13 @@ import { BaseToolOptions, BaseToolRunOptions, + StringToolOutput, Tool, - ToolInput, - ToolOutput + ToolInput } from 'bee-agent-framework/tools/base'; import { SchemaObject } from 'ajv'; import { Loaded } from '@mikro-orm/core'; +import { GetRunContext } from 'bee-agent-framework/context'; import { AgentContext } from '../execute.js'; @@ -41,26 +42,7 @@ export interface FunctionToolOptions extends BaseToolOptions { context: AgentContext; } -export class FunctionToolOutput extends ToolOutput { - constructor(public readonly output: string) { - super(); - } - - getTextContent(): string { - return this.output; - } - isEmpty(): boolean { - return false; - } - createSnapshot(): unknown { - return { - output: this.output - }; - } - loadSnapshot(_snapshot: unknown): void { - throw new Error('Method not implemented.'); - } -} +export class FunctionToolOutput extends StringToolOutput {} export class FunctionTool extends Tool { name: string; @@ -80,7 +62,11 @@ export class FunctionTool extends Tool return `run:${run.id}:call:${toolCallId}:output`; } - protected async _run(_: ToolInput, options: BaseToolRunOptions) { + protected async _run( + _: ToolInput, + __: BaseToolRunOptions | undefined, + run: GetRunContext + ) { const toolCall = this.options.context.toolCall; if (!(toolCall instanceof FunctionCall)) throw new Error('Invalid tool call'); @@ -114,8 +100,8 @@ export class FunctionTool extends Tool await ORM.em.flush(); resolve(new FunctionToolOutput(output)); }); - options.signal?.addEventListener('abort', () => { - reject(options.signal?.reason); + run.signal.addEventListener('abort', () => { + reject(run.signal.reason); }); }); } diff --git a/src/runs/execution/tools/helpers.ts b/src/runs/execution/tools/helpers.ts index 7143d0b..5ee0784 100644 --- a/src/runs/execution/tools/helpers.ts +++ b/src/runs/execution/tools/helpers.ts @@ -404,7 +404,7 @@ export async function finalizeToolCall( ); } else if (toolCall instanceof FunctionCall) { if (!(result instanceof FunctionToolOutput)) throw new TypeError(); - toolCall.output = result.output; + toolCall.output = result.result; } else if (toolCall instanceof FileSearchCall) { if (!(result instanceof FileSearchToolOutput)) throw new TypeError(); toolCall.results = result.results; diff --git a/src/runs/execution/tools/read-file-tool.ts b/src/runs/execution/tools/read-file-tool.ts index eea94b8..abef8b4 100644 --- a/src/runs/execution/tools/read-file-tool.ts +++ b/src/runs/execution/tools/read-file-tool.ts @@ -19,10 +19,12 @@ import { ToolInput, StringToolOutput, Tool, - ToolError + ToolError, + BaseToolRunOptions } from 'bee-agent-framework/tools/base'; import { z } from 'zod'; import { hasAtLeast } from 'remeda'; +import { GetRunContext } from 'bee-agent-framework/context'; import { File } from '@/files/entities/file.entity.js'; import { getExtractedText } from '@/files/extraction/helpers'; @@ -47,13 +49,17 @@ export class ReadFileTool extends Tool { }); } - protected async _run({ filename }: ToolInput): Promise { + protected async _run( + { filename }: ToolInput, + _: BaseToolRunOptions, + run: GetRunContext + ): Promise { const file = this.options.files.find((file) => file.filename === filename); if (!file) { throw new ToolError(`File ${filename} not found.`); } try { - const text = await getExtractedText(file); + const text = await getExtractedText(file, run.signal); if (text.length > this.options.fileSize) { throw new ToolError( `The text is too big (${text.length} characters). Maximum allowed size is ${this.options.fileSize} characters`, diff --git a/src/utils/datetime.ts b/src/utils/datetime.ts index 4bf9922..7d1352e 100644 --- a/src/utils/datetime.ts +++ b/src/utils/datetime.ts @@ -1,3 +1,19 @@ +/** + * 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 dayjs from 'dayjs'; import utc from 'dayjs/plugin/utc.js'; import timezone from 'dayjs/plugin/timezone.js';