diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ee3a6f..63abc9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,40 @@ # Changelog +## [0.0.20](https://github.com/i-am-bee/bee-api/compare/v0.0.19...v0.0.20) (2024-12-16) + +### Bug Fixes + +* **artifact:** update the share token only when the share property changes ([#139](https://github.com/i-am-bee/bee-api/issues/139)) ([c808d39](https://github.com/i-am-bee/bee-api/commit/c808d39b3f8c3af6ef206d553512ab7183fd440d)) + +## [0.0.19](https://github.com/i-am-bee/bee-api/compare/v0.0.18...v0.0.19) (2024-12-16) + +### Features + +* **user:** update default assistant for new users ([#138](https://github.com/i-am-bee/bee-api/issues/138)) ([233c138](https://github.com/i-am-bee/bee-api/commit/233c138f3075803bffeb127c1956c2742ee98ea5)) + +## [0.0.18](https://github.com/i-am-bee/bee-api/compare/v0.0.17...v0.0.18) (2024-12-15) + +### Features + +* **user:** add default assistant when user created ([#137](https://github.com/i-am-bee/bee-api/issues/137)) ([328c8de](https://github.com/i-am-bee/bee-api/commit/328c8ded996db1d8d8d649fbcff24c67f5cee2ed)) + +### Bug Fixes + +* **limits:** merge window and rate limits, fix streaming headers ([#134](https://github.com/i-am-bee/bee-api/issues/134)) ([e997180](https://github.com/i-am-bee/bee-api/commit/e9971803ae1f2e0d11f3a2d4cc0f4c194e83fb5a)) +* **readFile:** fix error for read file tool ([#133](https://github.com/i-am-bee/bee-api/issues/133)) ([05a59d9](https://github.com/i-am-bee/bee-api/commit/05a59d98b035bde280c71284b43c005a604d2128)) + +## [0.0.17](https://github.com/i-am-bee/bee-api/compare/v0.0.16...v0.0.17) (2024-12-13) + +### Features + +* **deps:** upgrade bee-agent-framework ([#135](https://github.com/i-am-bee/bee-api/issues/135)) ([ede714f](https://github.com/i-am-bee/bee-api/commit/ede714f7557bcb336f36385c74a6580c9e4d68f3)) + +## [0.0.16](https://github.com/i-am-bee/bee-api/compare/v0.0.15...v0.0.16) (2024-12-12) + +### Bug Fixes + +* **log:** update logged properties ([#130](https://github.com/i-am-bee/bee-api/issues/130)) ([e6ee1ec](https://github.com/i-am-bee/bee-api/commit/e6ee1ec5ed071b025b604a6410c5865279659f83)) + ## [0.0.15](https://github.com/i-am-bee/bee-api/compare/v0.0.14...v0.0.15) (2024-12-11) ### Features diff --git a/docs/add-native-tools.md b/docs/add-native-tools.md new file mode 100644 index 0000000..e950fd3 --- /dev/null +++ b/docs/add-native-tools.md @@ -0,0 +1,180 @@ +## Implement a Custom Native Typescript Tool (all the way to the UI) + + In this example we will implement the tool `RiddleTool`, all the way to the [Bee UI](https://github.com/i-am-bee/bee-ui). + +### Bee API + +### Create a new custom tool + +1. Create a new file in *src/runs/execution/tools* directory. In this example we will create *riddle-tool.ts*. +2. Copy and paste the [riddle example base tool](https://github.com/i-am-bee/bee-agent-framework/blob/main/examples/tools/custom/base.ts). In *riddle-tool.ts*. + +### Implement the tool in helpers and services + +#### Add the tool in *src/runs/execution/tools/helpers.ts* file. + +1. Import the tool in the file: + +```typescript +import { RiddleTool } from "./riddle-tool.js"; +``` + +2. Append the tool to the array `tools` in the `getTools` function: + +```typescript +export async function getTools(run: LoadedRun, context: AgentContext): Promise { + const tools: FrameworkTool[] = []; + + const vectorStores = getRunVectorStores(run.assistant.$, run.thread.$); + for (const vectorStore of vectorStores) { + vectorStore.lastActiveAt = new Date(); + } + + // Add the tool + const riddleUsage = run.tools.find( + (tool): tool is SystemUsage => + tool.type === ToolType.SYSTEM && tool.toolId === SystemTools.RIDDLE + ); + if (riddleUsage) { + tools.push(new RiddleTool()); + } + + ... +``` +3. Add the tool in `createToolCall` function: + +```typescript + ... + // Add the tool in the `else if` branch + } else if (tool instanceof RiddleTool) { + return new SystemCall({ + toolId: SystemTools.RIDDLE, + input: await tool.parse(input) + }); + } + throw new Error(`Unknown tool: ${tool.name}`); +} +``` +4. Add the tool in `finalizeToolCall` function: + +```typescript + ... + // Add the tool in the `switch` statement + case SystemTools.RIDDLE: { + // result can be an instance of arbitrary class + if (!(result instanceof StringToolOutput)) throw new TypeError(); + toolCall.output = result.result; + break; + } + } +```` + +#### Add the tool definition in *src/tools/entities/tool-calls/system-call.entity.ts* file: + +```typescript +export enum SystemTools { + WEB_SEARCH = 'web_search', + WIKIPEDIA = 'wikipedia', + WEATHER = 'weather', + ARXIV = 'arxiv', + READ_FILE = 'read_file', + RIDDLE = 'riddle', // Add the tool definition +} +``` + +#### Set the tool in the handlers in *src/tools/tools.service.ts* file: + +1. Import the tool in the file: + +```typescript +import { RiddleTool } from '@/runs/execution/tools/riddle-tool.js'; +``` + +2. Instance the tool in the `getSystemTools` function: + +```typescript +function getSystemTools() { + ... + + const systemTools = new Map(); + + const riddleTool = new RiddleTool(); // Add this line + + ... +``` +3. Set the tool at the end of `getSystemTools` function: + +```typescript + ... + // add this block of code + systemTools.set('riddle', { + type: ToolType.SYSTEM, + id: 'riddle', + createdAt: new Date('2024-11-22'), + ...riddleTool, + inputSchema: riddleTool.inputSchema.bind(riddleTool), + isExternal: false, // true if it accesses public internet + metadata: { + $ui_description_short: + 'It generates a random puzzle to test your knowledge.' + }, + userDescription: + 'It generates a random puzzle to test your knowledge.' + }); + + return systemTools; +} +``` +4. Add the tool in `listTools` function: + +```typescript + ... + + const systemTools: (SystemTool | undefined)[] = + !type || type.includes(ToolType.SYSTEM) + ? [ + allSystemTools.get(SystemTools.WEB_SEARCH), + allSystemTools.get(SystemTools.WIKIPEDIA), + allSystemTools.get(SystemTools.WEATHER), + allSystemTools.get(SystemTools.ARXIV), + allSystemTools.get('read_file'), + allSystemTools.get(SystemTools.RIDDLE) // add this line + ] + : []; + ... +``` +That's it! You have implemented the tool in the Bee API. :rocket: + +### Bee UI + +For the tool to be available in the UI, you need to follow these steps: + +1. Regenerate types (file *src/app/api/schema.d.ts* should change): + +```bash +pnpm schema:generate:api +``` + +2. Add the tool to `SYSTEM_TOOL_NAME` and `SYSTEM_TOOL_ICONS`in *src/modules/tools/hooks/useToolInfo.tsx* file: + +```typescript +const SYSTEM_TOOL_NAME: Record = { + wikipedia: 'Wikipedia', + web_search: 'WebSearch', + weather: 'OpenMeteo', + arxiv: 'Arxiv', + read_file: 'ReadFile', + riddle: 'Riddle', // Add this line +}; + +const SYSTEM_TOOL_ICONS: Record = { + wikipedia: Wikipedia, + web_search: IbmWatsonDiscovery, + weather: PartlyCloudy, + arxiv: Arxiv, + read_file: DocumentView, + riddle: Code, // Add this line +}; +``` + +That's it! You have implemented the tool in the Bee UI. :rocket: diff --git a/package.json b/package.json index 0fd44b2..4dbe084 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "bee-api", - "version": "0.0.15", + "version": "0.0.20", "license": "Apache-2.0", "author": "IBM Corp.", "type": "module", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 8479220..1a595c9 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -4134,8 +4134,8 @@ packages: resolution: {integrity: sha512-avsJQhyd+680gKXyG/sQc0nXaC6rBkPOfyHYcFb9+hdkqQkR9bdnkJ0AMZhke0oesPqIO+mFFJ+IdBc7mst4IA==} engines: {node: ^14.17.0 || ^16.13.0 || >=18.0.0} - nanoid@3.3.7: - resolution: {integrity: sha512-eSRppjcPIatRIMC1U6UngP8XFcz8MQWGQdt1MTBQ7NaAmvXDfvNxbvWV3x2y6CdEUciCSsDHDQZbhYaB8QEo2g==} + nanoid@3.3.8: + resolution: {integrity: sha512-WNLf5Sd8oZxOm+TzppcYk8gVOgP+l58xNy58D0nbUnOxOWRWvlcCV4kUF7ltmI6PsrLl/BgKEyS4mqsGChFN0w==} engines: {node: ^10 || ^12 || ^13.7 || ^14 || >=15.0.1} hasBin: true @@ -10080,7 +10080,7 @@ snapshots: mute-stream@1.0.0: {} - nanoid@3.3.7: {} + nanoid@3.3.8: {} napi-build-utils@1.0.2: {} @@ -10533,7 +10533,7 @@ snapshots: postcss@8.4.47: dependencies: - nanoid: 3.3.7 + nanoid: 3.3.8 picocolors: 1.1.0 source-map-js: 1.2.1 diff --git a/src/artifacts/artifacts.service.ts b/src/artifacts/artifacts.service.ts index f96111f..d5653d8 100644 --- a/src/artifacts/artifacts.service.ts +++ b/src/artifacts/artifacts.service.ts @@ -162,9 +162,9 @@ export async function updateArtifact({ } artifact.sourceCode = getUpdatedValue(source_code, artifact.sourceCode); } - if (shared === true) { + if (!artifact.accessToken && shared === true) { artifact.accessToken = getToken(); - } else if (shared === false) { + } else if (artifact.accessToken && shared === false) { artifact.accessToken = undefined; } await ORM.em.flush(); diff --git a/src/rate-limit.ts b/src/rate-limit.ts index 2104e30..12bff6f 100644 --- a/src/rate-limit.ts +++ b/src/rate-limit.ts @@ -36,6 +36,7 @@ const redis = createRedisClient({ export const rateLimitPlugin: FastifyPluginAsync = fp.default(async (app) => { await app.register(rateLimit, { global: true, + enableDraftSpec: true, max: (request: FastifyRequest) => { const authType = determineAuthType(request); switch (authType.type) { diff --git a/src/runs/execution/tools/read-file-tool.ts b/src/runs/execution/tools/read-file-tool.ts index b7f726f..d76665b 100644 --- a/src/runs/execution/tools/read-file-tool.ts +++ b/src/runs/execution/tools/read-file-tool.ts @@ -30,6 +30,7 @@ import { Emitter } from 'bee-agent-framework/emitter/emitter'; import { File } from '@/files/entities/file.entity.js'; import { getExtractedText } from '@/files/extraction/helpers'; +import { getJobLogger } from '@/logger'; export interface ReadFileToolOptions extends BaseToolOptions { fileSize: number; @@ -65,25 +66,28 @@ export class ReadFileTool extends Tool { if (!file) { throw new ToolError(`File ${filename} not found.`); } + let text: string; try { - 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`, - [], - { - isFatal: false, - isRetryable: true - } - ); - } + text = await getExtractedText(file, run.signal); + } catch (err) { + getJobLogger('runs').warn({ err }, 'Failed to get extracted text.'); - return new StringToolOutput('file content: \n' + text); - } catch { - throw new ToolError('This file is not a text file and can not be read.', [], { + throw new ToolError('Unable to read text from the file.', [], { isFatal: false, isRetryable: true }); } + 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`, + [], + { + isFatal: false, + isRetryable: true + } + ); + } + + return new StringToolOutput('file content: \n' + text); } } diff --git a/src/runs/runs.service.ts b/src/runs/runs.service.ts index 79f1b64..98a7a49 100644 --- a/src/runs/runs.service.ts +++ b/src/runs/runs.service.ts @@ -69,12 +69,14 @@ import { ensureRequestContextData } from '@/context.js'; 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'; export async function assertRunsQuota(newRuns = 1) { const count = await ORM.em.getRepository(Run).count({ createdBy: getProjectPrincipal(), createdAt: { $gte: getLatestDailyFixedTime().toDate() } }); + updateRateLimitHeadersWithDailyQuota({ quota: RUNS_QUOTA_DAILY, used: count }); if (count + newRuns > RUNS_QUOTA_DAILY) { throw new APIError({ message: 'Your daily runs quota has been exceeded', diff --git a/src/streaming/sse.ts b/src/streaming/sse.ts index decd85f..8f2e244 100644 --- a/src/streaming/sse.ts +++ b/src/streaming/sse.ts @@ -15,12 +15,17 @@ */ import { FastifyReply } from 'fastify'; +import { entries } from 'remeda'; import { Event } from './dtos/event.js'; export const init = (res: FastifyReply) => { res.hijack(); if (!res.raw.headersSent) { + const headers = res.getHeaders(); + entries(headers).forEach(([key, value]) => { + if (value) res.raw.setHeader(key, value); + }); res.raw.setHeader('Content-Type', 'text/event-stream'); res.raw.setHeader('Connection', 'keep-alive'); res.raw.setHeader('Cache-Control', 'no-cache,no-transform'); diff --git a/src/users/users.service.ts b/src/users/users.service.ts index e96158d..e63bcbf 100644 --- a/src/users/users.service.ts +++ b/src/users/users.service.ts @@ -34,6 +34,13 @@ import { OrganizationUserRole, ProjectRole } from '@/administration/entities/con import { Project } from '@/administration/entities/project.entity.js'; import { Organization } from '@/administration/entities/organization.entity.js'; import { IBM_ORGANIZATION_OWNER_ID } from '@/config.js'; +import { Assistant } from '@/assistants/assistant.entity.js'; +import { Agent, getDefaultModel } from '@/runs/execution/constants.js'; +import { SystemTools } from '@/tools/entities/tool-calls/system-call.entity.js'; +import { VECTOR_STORE_DEFAULT_MAX_NUM_RESULTS } from '@/vector-stores/constants.js'; +import { SystemUsage } from '@/tools/entities/tool-usages/system-usage.entity.js'; +import { FileSearchUsage } from '@/tools/entities/tool-usages/file-search-usage.entity.js'; +import { CodeInterpreterUsage } from '@/tools/entities/tool-usages/code-interpreter-usage.entity.js'; const getUserLogger = (userId: string) => getServiceLogger('user').child({ userId }); @@ -122,7 +129,32 @@ export async function createUser({ .getReference(organization.id, { wrapped: true }); user.defaultProject = ORM.em.getRepository(Project).getReference(project.id, { wrapped: true }); - await ORM.em.persistAndFlush([user, organization, orgUser, project, projectPrincipal]); + const assistant = new Assistant({ + model: getDefaultModel(), + agent: Agent.BEE, + tools: [ + new SystemUsage({ toolId: SystemTools.WEB_SEARCH }), + new SystemUsage({ toolId: SystemTools.READ_FILE }), + new FileSearchUsage({ maxNumResults: VECTOR_STORE_DEFAULT_MAX_NUM_RESULTS }), + new CodeInterpreterUsage() + ], + name: 'Bee', + project: ref(project), + createdBy: ref(projectPrincipal), + description: 'A general purpose agent for everyday tasks', + metadata: { + $ui_color: 'black', + $ui_icon: 'Bee', + '$ui_starterQuestion_c9f7253b-4fb3-4b2c-b576-d783f399ab6d': + 'Summarize key findings and methodology of the attached research paper.', + '$ui_starterQuestion_9207fe3c-9a01-4e37-bdf8-ddfbe9114397': + 'Bring me up to speed on the latest news and developments in the field of agentic AI. Be sure to cite your sources.', + '$ui_starterQuestion_919aa137-f90c-4580-b587-ab73d3d6b4b1': + 'Search for the top 5 most popular programming languages in 2024 and create a bar chart comparing the relative popularity of each language.' + } + }); + + await ORM.em.persistAndFlush([user, organization, orgUser, project, projectPrincipal, assistant]); getUserLogger(user.id).info({ externalId, metadata }, 'User created'); return toDto(user); } diff --git a/src/utils/rate-limit.ts b/src/utils/rate-limit.ts new file mode 100644 index 0000000..b9cf5e8 --- /dev/null +++ b/src/utils/rate-limit.ts @@ -0,0 +1,61 @@ +/** + * 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 { FastifyReply } from 'fastify'; + +import { dayjs, getLatestDailyFixedTime } from './datetime'; + +import { ensureRequestContextData } from '@/context'; + +function getNumericHeader(res: FastifyReply, header: string, fallback: number) { + const value = res.getHeader(header); + if (value === undefined) return fallback; + if (typeof value !== 'number') throw new Error('Invalid header type'); + return value; +} + +const RateLimitHeaders = { + LIMIT: 'ratelimit-limit', + REMAINING: 'ratelimit-remaining', + RESET: 'ratelimit-reset', + RETRY: 'retry-after' +} as const; + +export function updateRateLimitHeadersWithDailyQuota({ + quota, + used +}: { + quota: number; + used: number; +}) { + const res = ensureRequestContextData('res'); + res.header( + RateLimitHeaders.LIMIT, + Math.min(getNumericHeader(res, RateLimitHeaders.LIMIT, Infinity), quota) + ); + res.header( + RateLimitHeaders.REMAINING, + Math.min(getNumericHeader(res, RateLimitHeaders.REMAINING, Infinity), quota - used) + ); + if (quota === used) { + const reset = Math.max( + getNumericHeader(res, RateLimitHeaders.RESET, 0), + getLatestDailyFixedTime().add(1, 'day').unix() - dayjs().unix() + ); + res.header(RateLimitHeaders.RESET, reset); + res.header(RateLimitHeaders.RETRY, reset); + } +} diff --git a/src/vector-store-files/vector-store-files.service.ts b/src/vector-store-files/vector-store-files.service.ts index 121cd99..a7e3625 100644 --- a/src/vector-store-files/vector-store-files.service.ts +++ b/src/vector-store-files/vector-store-files.service.ts @@ -53,6 +53,7 @@ import { QueueName } from '@/jobs/constants.js'; import { getProjectPrincipal } from '@/administration/helpers.js'; import { VECTOR_STORE_FILE_QUOTA_DAILY } from '@/config.js'; import { dayjs, getLatestDailyFixedTime } from '@/utils/datetime.js'; +import { updateRateLimitHeadersWithDailyQuota } from '@/utils/rate-limit.js'; const getFileLogger = (vectorStoreFileIds?: string[]) => getServiceLogger('vector-store-files').child({ vectorStoreFileIds }); @@ -62,6 +63,7 @@ export async function assertVectorStoreFilesQuota(newFilesCount = 1) { createdBy: getProjectPrincipal(), createdAt: { $gte: getLatestDailyFixedTime().toDate() } }); + updateRateLimitHeadersWithDailyQuota({ quota: VECTOR_STORE_FILE_QUOTA_DAILY, used: count }); if (count + newFilesCount > VECTOR_STORE_FILE_QUOTA_DAILY) { throw new APIError({ message: 'Your daily vector store file quota has been exceeded', diff --git a/workers/python/poetry.lock b/workers/python/poetry.lock index 540da1c..aee6791 100644 --- a/workers/python/poetry.lock +++ b/workers/python/poetry.lock @@ -6255,4 +6255,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.0" python-versions = "~3.11" -content-hash = "b7c5999bf7bb5549653212da4a48285ca331b4aa8606c3cd7a5b9669e7ccf44d" +content-hash = "3481ba83425bf124a20d55ab7919488f419ebceb0f076b6b9b1ef91f36c4451f" diff --git a/workers/python/pyproject.toml b/workers/python/pyproject.toml index 359d9b5..f2c8943 100644 --- a/workers/python/pyproject.toml +++ b/workers/python/pyproject.toml @@ -19,7 +19,7 @@ opentelemetry-instrumentation-redis = "^0.49b2" optional = true [tool.poetry.group.docling.dependencies] -docling = "^2.9.0" +docling = "^2.10.0" [tool.poetry.group.unstructured] optional = true diff --git a/workers/python/python/config.py b/workers/python/python/config.py index a445f9a..3d2979e 100644 --- a/workers/python/python/config.py +++ b/workers/python/python/config.py @@ -52,5 +52,9 @@ def run_bullmq_workers(self) -> list[str]: otel_sdk_disabled: bool = False + docling_do_table_structure: bool = True + docling_pdf_do_ocr: bool = True + docling_advanced_chunker: bool = True + config = Config() diff --git a/workers/python/python/extraction/docling.py b/workers/python/python/extraction/docling.py index f5a844c..27c2b1d 100644 --- a/workers/python/python/extraction/docling.py +++ b/workers/python/python/extraction/docling.py @@ -17,8 +17,12 @@ import json import aioboto3 -from docling.document_converter import DocumentConverter +from docling.datamodel.base_models import InputFormat +from docling.document_converter import DocumentConverter, PdfFormatOption +from docling_core.transforms.chunker.hierarchical_chunker import HierarchicalChunker from docling_core.transforms.chunker.hybrid_chunker import HybridChunker +from docling.datamodel.pipeline_options import PdfPipelineOptions + from config import config from database import database @@ -27,6 +31,16 @@ S3_URL = f"s3://{config.s3_bucket_file_storage}" +converter = DocumentConverter(format_options={ + InputFormat.PDF: PdfFormatOption( + pipeline_options=PdfPipelineOptions( + do_table_structure=config.docling_do_table_structure, + do_ocr=config.docling_pdf_do_ocr) + ) +}) +chunker = HybridChunker( + tokenizer="BAAI/bge-small-en-v1.5") if config.docling_advanced_chunker else HierarchicalChunker() + async def docling_extraction(file): storage_id = file["storageId"] @@ -40,18 +54,17 @@ async def docling_extraction(file): aws_session_token=None, ) as s3: with tempfile.TemporaryDirectory() as tmp_dir: - # Use file_name to support file type discrimination. + # Use file_name to support file type discrimination. source_doc = f"{tmp_dir}/{file_name}" await s3.meta.client.download_file(config.s3_bucket_file_storage, storage_id, source_doc) - converter = DocumentConverter() result = await asyncio.to_thread(converter.convert, source_doc, max_num_pages=100, max_file_size=20971520) doc = result.document dict = doc.export_to_dict() markdown = doc.export_to_markdown() chunks = [{"text": c.text} - for c in list(HybridChunker(tokenizer="BAAI/bge-small-en-v1.5").chunk(doc))] + for c in list(await asyncio.to_thread(chunker.chunk, doc))] document_storage_id = f"{EXTRACTION_DIR}/{storage_id}/document.json" await s3.meta.client.put_object(