From 7ebb60a4206337b98c5dfd500c6e06a5df559106 Mon Sep 17 00:00:00 2001
From: Tomas Dvorak
Date: Fri, 23 Aug 2024 20:40:44 +0200
Subject: [PATCH] feat: init project
---
.env.template | 5 +
.github/ISSUE_TEMPLATE/bug_report.md | 31 +
.github/ISSUE_TEMPLATE/feature_request.md | 20 +
.gitignore | 139 +
.husky/pre-commit | 2 +
.lintstagedrc.json | 5 +
.nvmrc | 1 +
.release-it.json | 22 +
.yarnrc.yml | 1 +
CHANGELOG.md | 0
CODE_OF_CONDUCT.md | 76 +
LICENSE | 201 +
README.md | 173 +
docker-compose.yml | 60 +
docs/assets/Bee_Dark.svg | 15 +
docs/overview.md | 183 +
eslint.config.js | 69 +
examples/agents/bee.ts | 102 +
examples/agents/bee_reusable.ts | 34 +
examples/agents/simple.ts | 23 +
examples/helpers/io.ts | 73 +
examples/helpers/setup.ts | 4 +
examples/llms/chat.ts | 22 +
examples/llms/chatCallback.ts | 35 +
examples/llms/chatStream.ts | 23 +
examples/llms/providers/bam.ts | 32 +
examples/llms/providers/langchain.ts | 26 +
examples/llms/providers/ollama.ts | 42 +
examples/llms/providers/openai.ts | 20 +
examples/llms/providers/watsonx.ts | 53 +
examples/llms/structured.ts | 28 +
examples/llms/text.ts | 20 +
examples/template.ts | 68 +
package.json | 153 +
prettier.config.js | 8 +
scripts/copyright.sh | 34 +
.../bam/__snapshots__/llm.test.ts.snap | 833 ++
src/adapters/bam/chat.ts | 193 +
src/adapters/bam/chatPreset.ts | 109 +
src/adapters/bam/llm.test.ts | 163 +
src/adapters/bam/llm.ts | 369 +
src/adapters/langchain/llms/chat.ts | 225 +
src/adapters/langchain/llms/index.ts | 17 +
src/adapters/langchain/llms/llm.test.ts | 41 +
src/adapters/langchain/llms/llm.ts | 156 +
src/adapters/ollama/chat.ts | 215 +
src/adapters/ollama/llm.test.ts | 49 +
src/adapters/ollama/llm.ts | 190 +
src/adapters/ollama/shared.ts | 45 +
src/adapters/openai/chat.test.ts | 33 +
src/adapters/openai/chat.ts | 241 +
src/adapters/watsonx/chat.ts | 161 +
src/adapters/watsonx/llm.ts | 413 +
src/agents/base.ts | 78 +
.../bee/__snapshots__/parser.test.ts.snap | 26 +
src/agents/bee/agent.ts | 176 +
src/agents/bee/errors.ts | 19 +
src/agents/bee/parser.test.ts | 182 +
src/agents/bee/parser.ts | 293 +
src/agents/bee/prompts.ts | 117 +
src/agents/bee/runner.ts | 297 +
src/agents/bee/types.ts | 102 +
src/agents/manager.test.ts | 53 +
src/agents/manager.ts | 51 +
src/agents/types.ts | 43 +
src/cache/base.ts | 28 +
src/cache/decoratorCache.test.ts | 285 +
src/cache/decoratorCache.ts | 302 +
src/cache/fileCache.ts | 123 +
src/cache/nullCache.test.ts | 38 +
src/cache/nullCache.ts | 51 +
src/cache/slidingCache.test.ts | 73 +
src/cache/slidingCache.ts | 101 +
src/cache/unconstrainedCache.test.ts | 46 +
src/cache/unconstrainedCache.ts | 56 +
src/context.ts | 175 +
.../__snapshots__/emitter.test.ts.snap | 36 +
src/emitter/emitter.test.ts | 114 +
src/emitter/emitter.ts | 238 +
src/emitter/errors.ts | 19 +
src/emitter/types.ts | 56 +
src/emitter/utils.ts | 39 +
src/errors.test.ts | 95 +
src/errors.ts | 144 +
src/index.ts | 21 +
src/internals/env.ts | 40 +
src/internals/fetcher.ts | 144 +
.../__snapshots__/retryable.test.ts.snap | 164 +
src/internals/helpers/array.ts | 30 +
src/internals/helpers/cancellation.ts | 40 +
src/internals/helpers/counter.ts | 55 +
src/internals/helpers/general.ts | 35 +
src/internals/helpers/guards.ts | 28 +
src/internals/helpers/hash.ts | 29 +
src/internals/helpers/map.ts | 23 +
src/internals/helpers/number.ts | 21 +
src/internals/helpers/object.ts | 147 +
src/internals/helpers/promise.ts | 155 +
src/internals/helpers/prototype.ts | 35 +
src/internals/helpers/retry.ts | 72 +
src/internals/helpers/retryable.test.ts | 178 +
src/internals/helpers/retryable.ts | 218 +
src/internals/helpers/schema.ts | 87 +
src/internals/helpers/stream.ts | 26 +
src/internals/helpers/string.test.ts | 39 +
src/internals/helpers/string.ts | 100 +
src/internals/helpers/weakRef.ts | 71 +
src/internals/serializable.ts | 103 +
src/internals/types.ts | 112 +
src/llms/base.test.ts | 106 +
src/llms/base.ts | 249 +
src/llms/chat.ts | 110 +
src/llms/index.ts | 17 +
src/llms/llm.ts | 24 +
src/llms/primitives/message.ts | 63 +
src/llms/prompts.ts | 35 +
src/logger/logger.test.ts | 101 +
src/logger/logger.ts | 217 +
src/logger/pretty.js | 63 +
src/memory/base.ts | 79 +
src/memory/slidingWindowMemory.ts | 55 +
src/memory/summarizeMemory.ts | 109 +
src/memory/tokenMemory.test.ts | 60 +
src/memory/tokenMemory.ts | 128 +
src/memory/unconstrainedMemory.ts | 45 +
src/serializer/error.ts | 26 +
src/serializer/serializer.test.ts | 275 +
src/serializer/serializer.ts | 561 ++
src/serializer/utils.ts | 224 +
src/template.test.ts | 191 +
src/template.ts | 203 +
src/tools/arxiv.ts | 265 +
src/tools/base.test.ts | 516 +
src/tools/base.ts | 428 +
src/tools/custom.test.ts | 177 +
src/tools/custom.ts | 118 +
src/tools/llm.ts | 58 +
src/tools/python/output.ts | 68 +
src/tools/python/python.ts | 186 +
src/tools/python/storage.ts | 146 +
src/tools/search/base.ts | 50 +
src/tools/search/duckDuckGoSearch.test.ts | 121 +
src/tools/search/duckDuckGoSearch.ts | 155 +
src/tools/search/wikipedia.test.ts | 57 +
src/tools/search/wikipedia.ts | 314 +
src/tools/similarity.test.ts | 75 +
src/tools/similarity.ts | 97 +
src/tools/weather/openMeteo.ts | 173 +
src/tools/web/webCrawler.ts | 124 +
src/version.ts | 26 +
tests/e2e/adapters/langchain/chat.test.ts | 67 +
tests/e2e/adapters/sdk/chat.test.ts | 61 +
tests/e2e/adapters/watsonx/chat.test.ts | 101 +
tests/e2e/adapters/watsonx/llm.test.ts | 71 +
tests/e2e/agents/bee.test.ts | 137 +
tests/e2e/utils.ts | 134 +
tests/setup.ts | 26 +
tests/utils/file.ts | 32 +
tsconfig.json | 33 +
tsup.config.ts | 100 +
vitest.config.ts | 19 +
yarn.lock | 8778 +++++++++++++++++
162 files changed, 26339 insertions(+)
create mode 100644 .env.template
create mode 100644 .github/ISSUE_TEMPLATE/bug_report.md
create mode 100644 .github/ISSUE_TEMPLATE/feature_request.md
create mode 100644 .gitignore
create mode 100644 .husky/pre-commit
create mode 100644 .lintstagedrc.json
create mode 100644 .nvmrc
create mode 100644 .release-it.json
create mode 100644 .yarnrc.yml
create mode 100644 CHANGELOG.md
create mode 100644 CODE_OF_CONDUCT.md
create mode 100644 LICENSE
create mode 100644 README.md
create mode 100644 docker-compose.yml
create mode 100644 docs/assets/Bee_Dark.svg
create mode 100644 docs/overview.md
create mode 100644 eslint.config.js
create mode 100644 examples/agents/bee.ts
create mode 100644 examples/agents/bee_reusable.ts
create mode 100644 examples/agents/simple.ts
create mode 100644 examples/helpers/io.ts
create mode 100644 examples/helpers/setup.ts
create mode 100644 examples/llms/chat.ts
create mode 100644 examples/llms/chatCallback.ts
create mode 100644 examples/llms/chatStream.ts
create mode 100644 examples/llms/providers/bam.ts
create mode 100644 examples/llms/providers/langchain.ts
create mode 100644 examples/llms/providers/ollama.ts
create mode 100644 examples/llms/providers/openai.ts
create mode 100644 examples/llms/providers/watsonx.ts
create mode 100644 examples/llms/structured.ts
create mode 100644 examples/llms/text.ts
create mode 100644 examples/template.ts
create mode 100644 package.json
create mode 100644 prettier.config.js
create mode 100755 scripts/copyright.sh
create mode 100644 src/adapters/bam/__snapshots__/llm.test.ts.snap
create mode 100644 src/adapters/bam/chat.ts
create mode 100644 src/adapters/bam/chatPreset.ts
create mode 100644 src/adapters/bam/llm.test.ts
create mode 100644 src/adapters/bam/llm.ts
create mode 100644 src/adapters/langchain/llms/chat.ts
create mode 100644 src/adapters/langchain/llms/index.ts
create mode 100644 src/adapters/langchain/llms/llm.test.ts
create mode 100644 src/adapters/langchain/llms/llm.ts
create mode 100644 src/adapters/ollama/chat.ts
create mode 100644 src/adapters/ollama/llm.test.ts
create mode 100644 src/adapters/ollama/llm.ts
create mode 100644 src/adapters/ollama/shared.ts
create mode 100644 src/adapters/openai/chat.test.ts
create mode 100644 src/adapters/openai/chat.ts
create mode 100644 src/adapters/watsonx/chat.ts
create mode 100644 src/adapters/watsonx/llm.ts
create mode 100644 src/agents/base.ts
create mode 100644 src/agents/bee/__snapshots__/parser.test.ts.snap
create mode 100644 src/agents/bee/agent.ts
create mode 100644 src/agents/bee/errors.ts
create mode 100644 src/agents/bee/parser.test.ts
create mode 100644 src/agents/bee/parser.ts
create mode 100644 src/agents/bee/prompts.ts
create mode 100644 src/agents/bee/runner.ts
create mode 100644 src/agents/bee/types.ts
create mode 100644 src/agents/manager.test.ts
create mode 100644 src/agents/manager.ts
create mode 100644 src/agents/types.ts
create mode 100644 src/cache/base.ts
create mode 100644 src/cache/decoratorCache.test.ts
create mode 100644 src/cache/decoratorCache.ts
create mode 100644 src/cache/fileCache.ts
create mode 100644 src/cache/nullCache.test.ts
create mode 100644 src/cache/nullCache.ts
create mode 100644 src/cache/slidingCache.test.ts
create mode 100644 src/cache/slidingCache.ts
create mode 100644 src/cache/unconstrainedCache.test.ts
create mode 100644 src/cache/unconstrainedCache.ts
create mode 100644 src/context.ts
create mode 100644 src/emitter/__snapshots__/emitter.test.ts.snap
create mode 100644 src/emitter/emitter.test.ts
create mode 100644 src/emitter/emitter.ts
create mode 100644 src/emitter/errors.ts
create mode 100644 src/emitter/types.ts
create mode 100644 src/emitter/utils.ts
create mode 100644 src/errors.test.ts
create mode 100644 src/errors.ts
create mode 100644 src/index.ts
create mode 100644 src/internals/env.ts
create mode 100644 src/internals/fetcher.ts
create mode 100644 src/internals/helpers/__snapshots__/retryable.test.ts.snap
create mode 100644 src/internals/helpers/array.ts
create mode 100644 src/internals/helpers/cancellation.ts
create mode 100644 src/internals/helpers/counter.ts
create mode 100644 src/internals/helpers/general.ts
create mode 100644 src/internals/helpers/guards.ts
create mode 100644 src/internals/helpers/hash.ts
create mode 100644 src/internals/helpers/map.ts
create mode 100644 src/internals/helpers/number.ts
create mode 100644 src/internals/helpers/object.ts
create mode 100644 src/internals/helpers/promise.ts
create mode 100644 src/internals/helpers/prototype.ts
create mode 100644 src/internals/helpers/retry.ts
create mode 100644 src/internals/helpers/retryable.test.ts
create mode 100644 src/internals/helpers/retryable.ts
create mode 100644 src/internals/helpers/schema.ts
create mode 100644 src/internals/helpers/stream.ts
create mode 100644 src/internals/helpers/string.test.ts
create mode 100644 src/internals/helpers/string.ts
create mode 100644 src/internals/helpers/weakRef.ts
create mode 100644 src/internals/serializable.ts
create mode 100644 src/internals/types.ts
create mode 100644 src/llms/base.test.ts
create mode 100644 src/llms/base.ts
create mode 100644 src/llms/chat.ts
create mode 100644 src/llms/index.ts
create mode 100644 src/llms/llm.ts
create mode 100644 src/llms/primitives/message.ts
create mode 100644 src/llms/prompts.ts
create mode 100644 src/logger/logger.test.ts
create mode 100644 src/logger/logger.ts
create mode 100644 src/logger/pretty.js
create mode 100644 src/memory/base.ts
create mode 100644 src/memory/slidingWindowMemory.ts
create mode 100644 src/memory/summarizeMemory.ts
create mode 100644 src/memory/tokenMemory.test.ts
create mode 100644 src/memory/tokenMemory.ts
create mode 100644 src/memory/unconstrainedMemory.ts
create mode 100644 src/serializer/error.ts
create mode 100644 src/serializer/serializer.test.ts
create mode 100644 src/serializer/serializer.ts
create mode 100644 src/serializer/utils.ts
create mode 100644 src/template.test.ts
create mode 100644 src/template.ts
create mode 100644 src/tools/arxiv.ts
create mode 100644 src/tools/base.test.ts
create mode 100644 src/tools/base.ts
create mode 100644 src/tools/custom.test.ts
create mode 100644 src/tools/custom.ts
create mode 100644 src/tools/llm.ts
create mode 100644 src/tools/python/output.ts
create mode 100644 src/tools/python/python.ts
create mode 100644 src/tools/python/storage.ts
create mode 100644 src/tools/search/base.ts
create mode 100644 src/tools/search/duckDuckGoSearch.test.ts
create mode 100644 src/tools/search/duckDuckGoSearch.ts
create mode 100644 src/tools/search/wikipedia.test.ts
create mode 100644 src/tools/search/wikipedia.ts
create mode 100644 src/tools/similarity.test.ts
create mode 100644 src/tools/similarity.ts
create mode 100644 src/tools/weather/openMeteo.ts
create mode 100644 src/tools/web/webCrawler.ts
create mode 100644 src/version.ts
create mode 100644 tests/e2e/adapters/langchain/chat.test.ts
create mode 100644 tests/e2e/adapters/sdk/chat.test.ts
create mode 100644 tests/e2e/adapters/watsonx/chat.test.ts
create mode 100644 tests/e2e/adapters/watsonx/llm.test.ts
create mode 100644 tests/e2e/agents/bee.test.ts
create mode 100644 tests/e2e/utils.ts
create mode 100644 tests/setup.ts
create mode 100644 tests/utils/file.ts
create mode 100644 tsconfig.json
create mode 100644 tsup.config.ts
create mode 100644 vitest.config.ts
create mode 100644 yarn.lock
diff --git a/.env.template b/.env.template
new file mode 100644
index 00000000..6f1c9fbd
--- /dev/null
+++ b/.env.template
@@ -0,0 +1,5 @@
+# Optional
+CODE_INTERPRETER_URL=http://127.0.0.1:50051
+BEE_FRAMEWORK_LOG_PRETTY=true
+BEE_FRAMEWORK_LOG_LEVEL="info"
+BEE_FRAMEWORK_LOG_SINGLE_LINE="false"
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
new file mode 100644
index 00000000..df157ba8
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -0,0 +1,31 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: bug
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Steps to reproduce the behavior:
+1. Go to '...'
+2. Click on '....'
+3. Scroll down to '....'
+4. See error
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Screenshots / Code snippets**
+If applicable, add screenshots or code snippets to help explain your problem.
+
+**Set-up:**
+ - Bee version: [e.g. v0.0.3]
+ - Model provider [e.g. watsonx]
+
+**Additional context**
+Add any other context about the problem here.
diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md
new file mode 100644
index 00000000..bbcbbe7d
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/feature_request.md
@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 00000000..4cd5ca18
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,139 @@
+### Node template
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+lerna-debug.log*
+.pnpm-debug.log*
+
+# Diagnostic reports (https://nodejs.org/api/report.html)
+report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json
+
+# Runtime data
+pids
+*.pid
+*.seed
+*.pid.lock
+
+# Directory for instrumented libs generated by jscoverage/JSCover
+lib-cov
+
+# Coverage directory used by tools like istanbul
+coverage
+*.lcov
+
+# nyc test coverage
+.nyc_output
+
+# Grunt intermediate storage (https://gruntjs.com/creating-plugins#storing-task-files)
+.grunt
+
+# Bower dependency directory (https://bower.io/)
+bower_components
+
+# node-waf configuration
+.lock-wscript
+
+# Compiled binary addons (https://nodejs.org/api/addons.html)
+build/Release
+
+# Dependency directories
+node_modules/
+jspm_packages/
+
+# Snowpack dependency directory (https://snowpack.dev/)
+web_modules/
+
+# TypeScript cache
+*.tsbuildinfo
+
+# Optional npm cache directory
+.npm
+
+# Optional eslint cache
+.eslintcache
+
+# Optional stylelint cache
+.stylelintcache
+
+# Microbundle cache
+.rpt2_cache/
+.rts2_cache_cjs/
+.rts2_cache_es/
+.rts2_cache_umd/
+
+# Optional REPL history
+.node_repl_history
+
+# Output of 'npm pack'
+*.tgz
+
+# Yarn Integrity file
+.yarn-integrity
+
+# dotenv environment variable files
+.env
+.env.development.local
+.env.test.local
+.env.production.local
+.env.local
+
+# parcel-bundler cache (https://parceljs.org/)
+.cache
+.parcel-cache
+
+# Next.js build output
+.next
+out
+
+# Nuxt.js build / generate output
+.nuxt
+dist
+
+# Gatsby files
+.cache/
+# Comment in the public line in if your project uses Gatsby and not Next.js
+# https://nextjs.org/blog/next-9-1#public-directory-support
+# public
+
+# vuepress build output
+.vuepress/dist
+
+# vuepress v2.x temp and cache directory
+.temp
+.cache
+
+# Docusaurus cache and generated files
+.docusaurus
+
+# Serverless directories
+.serverless/
+
+# FuseBox cache
+.fusebox/
+
+# DynamoDB Local files
+.dynamodb/
+
+# TernJS port file
+.tern-port
+
+# Stores VSCode versions used for testing VSCode extensions
+.vscode-test
+
+# yarn v2
+.yarn/cache
+.yarn/unplugged
+.yarn/build-state.yml
+.yarn/install-state.gz
+.pnp.*
+
+.idea
+/experiments
+.npmrc
+
+# Code interpreter data
+/examples/tmp/code_interpreter/*
+/examples/tmp/local/*
diff --git a/.husky/pre-commit b/.husky/pre-commit
new file mode 100644
index 00000000..85e04b30
--- /dev/null
+++ b/.husky/pre-commit
@@ -0,0 +1,2 @@
+yarn lint-staged
+CI=true yarn lint && yarn run test:unit && yarn copyright
diff --git a/.lintstagedrc.json b/.lintstagedrc.json
new file mode 100644
index 00000000..325e2b51
--- /dev/null
+++ b/.lintstagedrc.json
@@ -0,0 +1,5 @@
+{
+ "*.{ts,js}": "eslint --fix",
+ "*.ts": "tsc-files --noEmit",
+ "*": "prettier --ignore-unknown --write"
+}
diff --git a/.nvmrc b/.nvmrc
new file mode 100644
index 00000000..b8e593f5
--- /dev/null
+++ b/.nvmrc
@@ -0,0 +1 @@
+20.15.1
diff --git a/.release-it.json b/.release-it.json
new file mode 100644
index 00000000..85efe4e1
--- /dev/null
+++ b/.release-it.json
@@ -0,0 +1,22 @@
+{
+ "plugins": {
+ "@release-it/conventional-changelog": {
+ "preset": {
+ "name": "conventionalcommits"
+ },
+ "header": "# Changelog",
+ "infile": "CHANGELOG.md"
+ }
+ },
+ "npm": {
+ "skipChecks": true,
+ "publish": true
+ },
+ "hooks": {
+ "before:init": ["yarn lint", "yarn ts:check", "yarn test:all"],
+ "after:bump": ["yarn build"]
+ },
+ "github": {
+ "release": true
+ }
+}
diff --git a/.yarnrc.yml b/.yarnrc.yml
new file mode 100644
index 00000000..3186f3f0
--- /dev/null
+++ b/.yarnrc.yml
@@ -0,0 +1 @@
+nodeLinker: node-modules
diff --git a/CHANGELOG.md b/CHANGELOG.md
new file mode 100644
index 00000000..e69de29b
diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md
new file mode 100644
index 00000000..f3e51c64
--- /dev/null
+++ b/CODE_OF_CONDUCT.md
@@ -0,0 +1,76 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and expression,
+level of experience, education, socio-economic status, nationality, personal
+appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+ advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+ address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+ professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies both within project spaces and in public spaces
+when an individual is representing the project or its community. Examples of
+representing a project or community include using an official project e-mail
+address, posting via an official social media account, or acting as an appointed
+representative at an online or offline event. Representation of a project may be
+further defined and clarified by project maintainers.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at [IBM Cloud Support](https://www.ibm.com/cloud/support).
+All complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see
+https://www.contributor-covenant.org/faq
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 00000000..261eeb9e
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,201 @@
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ 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.
diff --git a/README.md b/README.md
new file mode 100644
index 00000000..ff5bd021
--- /dev/null
+++ b/README.md
@@ -0,0 +1,173 @@
+
+
+
Bee Agent Framework
+
+
+
+
+
+
+
+
+Open-source framework for building, deploying, and serving powerful agentic workflows at scale.
+
+The Bee framework makes it easy to build agentic worfklows with leading proprietary and open-source models. Weβre working on bringing model-agnostic support to any LLM to help developers avoid model provider lock-in and embrace the latest open-source LLMs.
+
+## Key Features
+
+- π€ **AI agents**: Use our powerful Bee agent or build your own.
+- π οΈ **Tools**: Use our built-in tools and create your own in Javascript or Python.
+- π©βπ» **Code interpreter**: Run code safely in a sandbox container.
+- πΎ **Memory**: Multiple strategies to optimize token spend.
+- βΈοΈ **Serialization** Handle complex agentic workflows and easily pause/resume them without losing state.
+- π **Traceability**: Get full visibility of your agentβs inner workings [link to docs page], log all running events and use our MLflow integration to debug performance [link to instructions to setup MLflow].
+- ποΈ **Production-level** control with caching and error handling [link to docs]
+- π§ (Coming soon) **Evaluation**: Run evaluation jobs with your own data source (custom csv or Airtable data).
+- π§ (Coming soon) **Model-agnostic support**: Change model providers in 1 line of code without breaking your agentβs functionality.
+- π§ (Coming soon) **Chat UI**: Serve your agent to users in a delightful GUI with built-in transparency, explainability, and user controls.
+- ... more on our [Roadmap](#roadmap)
+
+## Get started with Bee
+
+### Installation
+
+```shell
+npm install bee-agent-framework
+```
+
+```shell
+yarn add bee-agent-framework
+```
+
+#### Example
+
+```typescript
+import { BeeAgent } from "bee-agent-framework/agents/bee/agent.js";
+import { OllamaChatLLM } from "bee-agent-framework/adapters/ollama/chat.js";
+import { TokenMemory } from "bee-agent-framework/memory/tokenMemory.js";
+import { DuckDuckGoSearchTool } from "bee-agent-framework/tools/search/duckDuckGoSearch.js";
+import { OpenMeteoTool } from "bee-agent-framework/tools/weather/openMeteo.js";
+
+const llm = new OllamaChatLLM(); // default is llama3.1 (7b), it is recommended to use 70b model
+const agent = new BeeAgent({
+ llm, // for more explore 'bee-agent-framework/adapters'
+ memory: new TokenMemory({ llm }), // for more explore 'bee-agent-framework/memory'
+ tools: [new DuckDuckGoSearchTool(), new OpenMeteoTool()], // for more explore 'bee-agent-framework/tools'
+});
+
+const response = await agent
+ .run({ prompt: "What's the current weather in Las Vegas?" })
+ .observe((emitter) => {
+ emitter.on("update", async ({ data, update, meta }) => {
+ console.log(`Agent (${update.key}) π€ : `, update.value);
+ });
+ });
+
+console.log(`Agent π€ : `, response.result.text);
+```
+
+β‘οΈ See a more [advanced example](./examples/agents/bee.ts).
+
+β‘οΈ All examples can be found in the [examples](./examples) directory.
+
+### Local Installation (Python Interpreter + Interactive CLI)
+
+> _Note: `yarn` should be installed via Corepack ([tutorial](https://yarnpkg.com/corepack))_
+
+> _Note: To make any asset available to a local code interpreter place them the following directory: ./examples/tmp/local_
+
+> _Note: Docker distribution with support for compose is required, the following are supported:_
+>
+> - [Rancher](https://www.rancher.com/) - recommended
+> - [Docker](https://www.docker.com/)
+> - [Podman](https://podman.io/) - requires [compose](https://podman-desktop.io/docs/compose/setting-up-compose) and **rootful machine** (if your current machine is rootless, please create a new one)
+
+1. Clone the repository `git clone git@github.com:i-am-bee/bee-agent-framework`.
+2. Install dependencies `yarn install`.
+3. Create `.env` (from `.env.template`) and fill in missing values (if any).
+4. Start the code interpreter `yarn run infra:start-code-interpreter`.
+5. Start the agent `yarn run start:bee` (it runs ./examples/agents/bee.ts file).
+
+
+### π οΈ Tools
+
+| Name | Description |
+| ------------------------------------------------------------------- | ------------------------------------------------------------------------- |
+| `PythonTool` | Run arbitrary Python code in the remote environment. |
+| `WikipediaTool` | Search for data on Wikipedia. |
+| `DuckDuckGoTool` | Search for data on DuckDuckGo. |
+| `LLMTool` | Uses an LLM to process input data. |
+| `DynamicTool` | Construct to create dynamic tools. |
+| `ArXivTool` | Retrieves research articles published on arXiv. |
+| `WebCrawlerTool` | Retrieves content of an arbitrary website. |
+| `CustomTool` | Runs your own Python function in the remote environment. |
+| `OpenMeteoTool` | Retrieves current, previous, or upcoming weather for a given destination. |
+| β [Request](https://github.com/i-am-bee/bee-agent-framework/discussions) | |
+
+### ποΈ Adapters (LLM - Inference providers)
+
+| Name | Description |
+|--------------------------------------------------------------------------|-----------------------------------------------------------------------------------------|
+| `Ollama` | LLM + ChatLLM support ([example](./examples/llms/providers/ollama.ts)) |
+| `LangChain` | Use any LLM that LangChain supports ([example](./examples/llms/providers/langchain.ts)) |
+| `WatsonX` | LLM + ChatLLM support ([example](./examples/llms/providers/watsonx.ts)) |
+| `BAM (IBM Internal)` | LLM + ChatLLM support ([example](./examples/llms/providers/bam.ts)) |
+| β [Request](https://github.com/i-am-bee/bee-agent-framework/discussions) | |
+
+
+### π¦ Modules
+
+The source directory (`src`) provides numerous modules that one can use.
+
+| Name | Description |
+| -------------- |--------------------------------------------------------------------------------------------|
+| **agents** | Base classes defining the common interface for agent. |
+| **llms** | Base classes defining the common interface for text inference (standard or chat). |
+| **template** | Prompt Templating system based on `Mustache` with various improvements_. |
+| **memory** | Various types of memories to use with agent. |
+| **tools** | Tools that an agent can use. |
+| **cache** | Preset of different caching approaches that can be used together with tools. |
+| **errors** | Base framework error classes used by each module. |
+| **adapters** | Concrete implementations of given modules for different environments. |
+| **logger** | Core component for logging all actions within the framework. |
+| **serializer** | Core component for the ability to serialize/deserialize modules into the serialized format. |
+| **version** | Constants representing the framework (e.g., latest version) |
+| **internals** | Modules used by other modules within the framework. |
+
+To see more in-depth explanation see (docs)(./docs/overview.md).
+
+## Tutorials
+
+π§ Coming soon π§
+
+## Roadmap
+
+- Evaluation with MLFlow integration
+- JSON encoder/decoder for model-agnostic support
+- Chat Client (GUI)
+- Structured outputs
+- Improvements to base Bee agent
+- Guardrails
+- π§ TBD π§
+
+## Contribution guidelines
+
+The Bee Agent Framework is an open-source project and we β€οΈ contributions.
+
+## Feature contributions
+
+You can get started with any ticket market as βgood first issueβ.
+
+Have an idea for a new feature? We recommend you first talk to a maintainer prior to spending a lot of time making a pull request that may not align with the project roadmap.
+
+## Bugs
+
+We are using [GitHub Issues](https://github.com/i-am-bee/bee-agent-framework/issues) to manage our public bugs. We keep a close eye on this, so before filing a new issue, please check to make sure it hasn't already been logged.
+
+## Code of conduct
+
+This project and everyone participating in it are governed by the [Code of Conduct](./CODE_OF_CONDUCT.md). By participating, you are expected to uphold this code. Please read the [full text](./CODE_OF_CONDUCT.md) so that you can read which actions may or may not be tolerated.
+
+## Legal notice
+
+All content in these repositories including code has been provided by IBM under the associated open source software license and IBM is under no obligation to provide enhancements, updates, or support. IBM developers produced this code as an open source project (not as an IBM product), and IBM makes no assertions as to the level of quality nor security, and will not be maintaining this code going forward.
diff --git a/docker-compose.yml b/docker-compose.yml
new file mode 100644
index 00000000..8ac7286d
--- /dev/null
+++ b/docker-compose.yml
@@ -0,0 +1,60 @@
+services:
+ bee-code-interpreter:
+ image: "iambeeagent/bee-code-interpreter:latest"
+ environment:
+ KUBECONFIG: /root/kubeconfig.yaml
+ APP_EXECUTOR_IMAGE: "iambeeagent/bee-code-executor:latest"
+ APP_EXECUTOR_CONTAINER_RESOURCES: "{}"
+ APP_FILE_STORAGE_PATH: "/storage"
+ volumes:
+ - ${CODE_INTEPRETER_TMPDIR:-./examples/tmp/code_interpreter}:/storage
+ - k3s-kubeconfig:/kube
+ ports:
+ - "50051:50051"
+ entrypoint: ["/bin/sh", "-c"]
+ command:
+ - >
+ echo 'Updating kubeconfig' &&
+ while [ ! -f $${KUBECONFIG} ] || ! kubectl get namespace default 2>/dev/null 1>&2; do
+ sleep 1;
+ sed 's|127.0.0.1|bee-code-interpreter-k3s|g' /kube/kubeconfig.yaml >$${KUBECONFIG} 2>/dev/null;
+ done &&
+ echo 'Kubeconfig updated successfully' &&
+
+ echo "Pulling executor image $${APP_EXECUTOR_IMAGE}" &&
+ if kubectl get job/pull-image 2>/dev/null 1>&2; then
+ echo "pull-image job already exists";
+ else
+ kubectl create job pull-image --image "$${APP_EXECUTOR_IMAGE}" -- /bin/sh -c "echo done";
+ fi &&
+ kubectl wait --for=condition=complete job/pull-image --timeout 3600s &&
+ echo 'Image pulled successfully' &&
+
+ python -m code_interpreter
+ depends_on:
+ - bee-code-interpreter-k3s
+
+ bee-code-interpreter-k3s:
+ image: "rancher/k3s:v1.29.7-k3s1"
+ command: ["server", "--tls-san", "bee-code-interpreter-k3s"]
+ tmpfs:
+ - /run
+ - /var/run
+ ulimits:
+ nproc: 65535
+ nofile:
+ soft: 65535
+ hard: 65535
+ privileged: true
+ restart: always
+ environment:
+ - K3S_TOKEN=secret-token
+ - K3S_KUBECONFIG_OUTPUT=/output/kubeconfig.yaml
+ - K3S_KUBECONFIG_MODE=644
+ volumes:
+ - k3s-containerd:/var/lib/rancher/k3s/agent/containerd
+ - k3s-kubeconfig:/output
+
+volumes:
+ k3s-kubeconfig:
+ k3s-containerd:
diff --git a/docs/assets/Bee_Dark.svg b/docs/assets/Bee_Dark.svg
new file mode 100644
index 00000000..d80a4083
--- /dev/null
+++ b/docs/assets/Bee_Dark.svg
@@ -0,0 +1,15 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/docs/overview.md b/docs/overview.md
new file mode 100644
index 00000000..3677473f
--- /dev/null
+++ b/docs/overview.md
@@ -0,0 +1,183 @@
+# Overview
+
+## π¦ Modules
+
+The source directory (`src`) provides numerous modules that one can use.
+
+| Name | Description |
+| -------------- |--------------------------------------------------------------------------------------------|
+| **agents** | Base classes defining the common interface for agent. |
+| **llms** | Base classes defining the common interface for text inference (standard or chat). |
+| **template** | Prompt Templating system based on `Mustache` with various improvements_. |
+| **memory** | Various types of memories to use with agent. |
+| **tools** | Tools that an agent can use. |
+| **cache** | Preset of different caching approaches that can be used together with tools. |
+| **errors** | Base framework error classes used by each module. |
+| **adapters** | Concrete implementations of given modules for different environments. |
+| **logger** | Core component for logging all actions within the framework. |
+| **serializer** | Core component for the ability to serialize/deserialize modules into the serialized format. |
+| **version** | Constants representing the framework (e.g., latest version) |
+| **internals** | Modules used by other modules within the framework. |
+
+### Emitter
+
+> Location within the framework `bee-agent-framework/emitter`.
+
+An emitter is a core functionality of the framework that gives you the ability to see what is happening under the hood. Because the vast majority of the frameworks components uses it, you can easily observe them.
+
+```typescript
+import { Emitter, EventMeta } from "@/emitter/emitter.js";
+
+Emitter.root.match("*.*", (data: unknown, event: EventMeta) => {
+ console.info(event.path, { data })
+})
+```
+
+π§ TBD
+
+**π [Emitter + Agent example](../examples/agents/bee.ts) π**
+
+### LLMs
+
+A Large Language Model (LLM) is an AI designed to understand and generate human-like text.
+Trained on extensive text data, LLMs learn language patterns, grammar, context, and basic reasoning to perform tasks like text completion, translation, summarization, and answering questions.
+
+The framework defines LLM via the `BaseLLM` class (`src/llms/base.ts`).
+Typically, you either want to implement `LLM` (string <-> string) or `ChatLLM` (Message <-> Message).
+
+**π [See Examples](../examples/llms) π**
+
+### Templates
+
+> Location within the framework `bee-agent-framework/template`.
+
+Prompt templating involves creating structured inputs to guide LLMs in generating desired responses.
+Users can effectively steer the model's output towards particular formats or tones by designing specific templates.
+This technique enhances the model's ability to produce consistent and relevant results, making it particularly useful for applications like automated content creation, customer service scripts, and interactive storytelling.
+Prompt templating ensures that the LLM's responses align more closely with the intended context and purpose.
+
+**π [See Examples](../examples/template.ts) π**
+
+### Agents
+
+π§ TBD π§
+
+**How to implement your agent runtime?**
+
+By implementing the agent base interface defined in `src/agents/base.ts.`
+
+### Memory
+
+> Location within the framework `bee-agent-framework/memory`.
+
+Memory in the context of an agent refers to the system's capability to store, recall, and utilize information from past interactions.
+This enables the agent to maintain context over time, improve its responses based on previous exchanges, and provide a more personalized experience.
+
+The framework provides out-of-the-box following types of memories.
+
+| Name | Description |
+| --------------------- | -------------------------------------------------------------------------------------------------------------- |
+| `UnconstrainedMemory` | Unlimited size. It is suitable if your context window is huge. |
+| `SlidingWindowMemory` | Keeps last `k` messages in the memory. The oldest ones are deleted. |
+| `TokenMemory` | Ensures that the number of tokens of all messages is below the given threshold. The oldest are removed. |
+| `SummarizeMemory` | Only a single summarization of the conversation is preserved. Summarization is updated with every new message. |
+| β [Request](https://github.com/i-am-bee/bee-agent-framework/discussions) | |
+
+### Tools
+
+> Location within the framework `bee-agent-framework/tools`.
+
+Tools in the context of an agent refer to additional functionalities or capabilities integrated with the agent to perform specific tasks beyond text processing.
+These tools extend the agent's abilities, allowing it to interact with external systems, access information, and execute actions.
+
+The framework provides out-of-the-box tools.
+
+| Name | Description |
+| ------------------------------------------------------------------- | ------------------------------------------------------------------------- |
+| `PythonTool` | Run arbitrary Python code in the remote environment. |
+| `WikipediaTool` | Search for data on Wikipedia. |
+| `DuckDuckGoTool` | Search for data on DuckDuckGo. |
+| `LLMTool` | Uses an LLM to process input data. |
+| `DynamicTool` | Construct to create dynamic tools. |
+| `ArXivTool` | Retrieves research articles published on arXiv. |
+| `WebCrawlerTool` | Retrieves content of an arbitrary website. |
+| `CustomTool` | Runs your own Python function in the remote environment. |
+| `OpenMeteoTool` | Retrieves current, previous, or upcoming weather for a given destination. |
+| β [Request](https://github.com/i-am-bee/bee-agent-framework/discussions) | |
+
+To create your own tool, you need to either implement the `BaseTool` class ([example](../examples/tools/customHttpRequest.ts)) or use `DynamicTool.`
+
+### Cache
+
+> Location within the framework `bee-agent-framework/cache`.
+
+> Note: Cache can be used directly with Tools. Pass the appropriate `Cache` instance to the `Tool` constructor.
+
+Caching is a process used to temporarily store copies of data or computations in a cache (a storage location) to facilitate faster access upon future requests.
+The primary purpose of caching is to improve the efficiency and performance of systems by reducing the need to repeatedly fetch or compute the same data from a slower or more resource-intensive source.
+
+The framework provides out-of-the-box following cache implementations.
+
+| Name | Description |
+|----------------------|--------------------------------------------------------------------|
+| `UnconstrainedCache` | Unlimited size. |
+| `FileCache` | Saves/Loads entries to/from a file. |
+| `SlidingCache` | Keeps last `k` entries in the memory. The oldest ones are deleted. |
+| `NullCache` | Disables caching. |
+| β [Request](https://github.com/i-am-bee/bee-agent-framework/discussions) | |
+
+To create your cache implementation, you must implement the `BaseCache` class.
+
+### Errors
+
+> Location within the framework `bee-agent-framework/error`.
+
+> Note: We guarantee that every framework-related error is an instance of the `FrameworkError` class.
+
+π§ TBD
+
+### Logger
+
+> Location within the framework `bee-agent-framework/logger`.
+> To log all events in the framework set log level to 'TRACE' (root logger observes the root emitter).
+
+A logger is a component used for recording and tracking events, errors, and other significant actions that occur during an application's execution.
+The primary purpose of a logger is to provide developers and system administrators with insights into the application's behavior, performance, and potential issues.
+
+Every component within the framework uses the `Logger` class either by directly creating an instance of it or because it is being passed from the creator.
+
+```typescript
+import { Logger, LoggerLevel } from "bee-agent-framework/logger";
+
+Logger.defaults.pretty = true; // (override default settings)
+const root = Logger.root; // get root logger
+root.level = LogerLevel.WARN; // update the logger level (default is LoggerLevel.INFO)
+```
+
+Some of the `Logger` defaults can be controlled via the following environmental variables.
+
+- `BEE_FRAMEWORK_LOG_LEVEL`
+- `BEE_FRAMEWORK_LOG_PRETTY`
+
+> Note: The framework `Logger` class is an abstraction on top of the most popular `pino` logger.
+
+### Serializer
+
+> Location within the framework `bee-agent-framework/serializer`.
+
+Serialization is the process of converting complex data structures or objects into a format that can be easily stored, transmitted, and reconstructed later.
+
+Most parts of the framework implement the internal `Serializable` class, which exposes the following functionalities.
+
+- `createSnapshot` (method)
+- `loadSnapshot` (method)
+
+- `fromSerialized` (static method)
+- `fromSnapshot` (static method)
+
+If you want to serialize something that the framework doesn't know how to process, the following error will be thrown: `SerializerError`.
+To resolve such an issue, you need to tell (register) the appropriate class to the framework via the `Serializer.register` method.
+
+**π [Emitter + Agent example](../examples/agents/bee_reusable.ts) π**
+
+## π§ More content TBD π§
diff --git a/eslint.config.js b/eslint.config.js
new file mode 100644
index 00000000..eab964ba
--- /dev/null
+++ b/eslint.config.js
@@ -0,0 +1,69 @@
+// @ts-check
+
+import eslint from "@eslint/js";
+import tseslint from "typescript-eslint";
+import prettierConfig from "eslint-config-prettier";
+import unusedImports from "eslint-plugin-unused-imports";
+
+export default tseslint.config(
+ eslint.configs.recommended,
+ ...tseslint.configs.strict,
+ ...tseslint.configs.stylistic,
+ {
+ ignores: ["**/*.js"],
+ languageOptions: {
+ parserOptions: {
+ project: "./tsconfig.json",
+ },
+ },
+ plugins: {
+ "unused-imports": unusedImports,
+ },
+ rules: {
+ "no-restricted-imports": [
+ "error",
+ {
+ patterns: [
+ {
+ group: ["../", "src/"],
+ message: "Relative imports are not allowed.",
+ },
+ ],
+ },
+ ],
+ "@typescript-eslint/no-explicit-any": "off",
+ "@typescript-eslint/no-unused-vars": "off",
+ "@typescript-eslint/no-empty-function": "off",
+ "@typescript-eslint/no-extraneous-class": "off",
+ "@typescript-eslint/no-empty-interface": "off",
+ "@typescript-eslint/no-non-null-assertion": "off",
+ "@typescript-eslint/no-floating-promises": "error",
+ "unused-imports/no-unused-imports": "error",
+ "unused-imports/no-unused-vars": [
+ "warn",
+ {
+ vars: "all",
+ varsIgnorePattern: "^_",
+ args: "after-used",
+ argsIgnorePattern: "^_",
+ },
+ ],
+ "@typescript-eslint/ban-ts-comment": "off",
+ "@typescript-eslint/no-empty-object-type": "off",
+ },
+ },
+ {
+ files: ["examples/**"],
+ rules: {
+ "no-restricted-imports": "off",
+ "@typescript-eslint/no-unused-vars": "off",
+ "unused-imports/no-unused-vars": "off",
+ },
+ },
+ prettierConfig,
+ {
+ rules: {
+ curly: ["error", "all"],
+ },
+ },
+);
diff --git a/examples/agents/bee.ts b/examples/agents/bee.ts
new file mode 100644
index 00000000..bfc83a30
--- /dev/null
+++ b/examples/agents/bee.ts
@@ -0,0 +1,102 @@
+import "dotenv/config.js";
+import { BeeAgent } from "@/agents/bee/agent.js";
+import { createConsoleReader } from "../helpers/io.js";
+import { FrameworkError } from "@/errors.js";
+import { TokenMemory } from "@/memory/tokenMemory.js";
+import { Logger } from "@/logger/logger.js";
+import { PythonTool } from "@/tools/python/python.js";
+import { LocalPythonStorage } from "@/tools/python/storage.js";
+import { DuckDuckGoSearchTool } from "@/tools/search/duckDuckGoSearch.js";
+import { WikipediaTool } from "@/tools/search/wikipedia.js";
+import { OpenMeteoTool } from "@/tools/weather/openMeteo.js";
+import { dirname } from "node:path";
+import { fileURLToPath } from "node:url";
+import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
+
+Logger.root.level = "silent"; // disable internal logs
+const logger = new Logger({ name: "app", level: "trace" });
+
+const llm = new OllamaChatLLM({
+ modelId: "llama3.1", // llama3.1:70b for better performance
+});
+
+const codeInterpreterUrl = process.env.CODE_INTERPRETER_URL;
+const __dirname = dirname(fileURLToPath(import.meta.url));
+
+const agent = new BeeAgent({
+ llm,
+ memory: new TokenMemory({ llm }),
+ tools: [
+ new DuckDuckGoSearchTool(),
+ // new WebCrawlerTool(), // HTML web page crawler
+ new WikipediaTool(),
+ new OpenMeteoTool(), // weather tool
+ // new ArXivTool(), // research papers
+ // new DynamicTool() // custom python tool
+ ...(codeInterpreterUrl
+ ? [
+ new PythonTool({
+ codeInterpreter: { url: codeInterpreterUrl },
+ storage: new LocalPythonStorage({
+ interpreterWorkingDir: `${__dirname}/tmp/code_interpreter`,
+ localWorkingDir: `${__dirname}/tmp/local`,
+ }),
+ }),
+ ]
+ : []),
+ ],
+});
+
+const reader = createConsoleReader();
+if (codeInterpreterUrl) {
+ reader.write("π οΈ System", "Please ensure that the code interpreter is running.");
+}
+
+try {
+ for await (const { prompt } of reader) {
+ const response = await agent
+ .run(
+ { prompt },
+ {
+ execution: {
+ maxRetriesPerStep: 3,
+ totalMaxRetries: 10,
+ maxIterations: 20,
+ },
+ },
+ )
+ .observe((emitter) => {
+ // emitter.on("start", () => {
+ // reader.write(`Agent π€ : `, "starting new iteration");
+ // });
+ emitter.on("error", ({ error }) => {
+ reader.write(`Agent π€ : `, FrameworkError.ensure(error).dump());
+ });
+ emitter.on("retry", () => {
+ reader.write(`Agent π€ : `, "retrying the action...");
+ });
+ emitter.on("update", async ({ data, update, meta }) => {
+ // log 'data' to see the whole state
+ // to log only valid runs (no errors), check if meta.success === true
+ reader.write(`Agent (${update.key}) π€ : `, update.value);
+ });
+ emitter.on("partialUpdate", ({ data, update, meta }) => {
+ // ideal for streaming (line by line)
+ // log 'data' to see the whole state
+ // to log only valid runs (no errors), check if meta.success === true
+ // reader.write(`Agent (partial ${update.key}) π€ : `, update.value);
+ });
+
+ // To observe all events (uncomment following block)
+ // emitter.match("*.*", async (data: unknown, event) => {
+ // logger.trace(event, `Received event "${event.path}"`);
+ // });
+ });
+
+ reader.write(`Agent π€ : `, response.result.text);
+ }
+} catch (error) {
+ logger.error(FrameworkError.ensure(error).dump());
+} finally {
+ process.exit(0);
+}
diff --git a/examples/agents/bee_reusable.ts b/examples/agents/bee_reusable.ts
new file mode 100644
index 00000000..2f315d03
--- /dev/null
+++ b/examples/agents/bee_reusable.ts
@@ -0,0 +1,34 @@
+import "dotenv/config.js";
+import { BeeAgent } from "@/agents/bee/agent.js";
+import { DuckDuckGoSearchTool } from "@/tools/search/duckDuckGoSearch.js";
+import { UnconstrainedMemory } from "@/memory/unconstrainedMemory.js";
+import { OpenAIChatLLM } from "@/adapters/openai/chat.js";
+
+// We create an agent
+let agent = new BeeAgent({
+ llm: new OpenAIChatLLM(),
+ tools: [new DuckDuckGoSearchTool()],
+ memory: new UnconstrainedMemory(),
+});
+
+// We ask the agent
+let prompt = "Who is the president of USA?";
+console.info(prompt);
+const response = await agent.run({
+ prompt,
+});
+console.info(response.result.text);
+
+// We can save (serialize) the agent
+const json = agent.serialize();
+
+// We reinitialize the agent to the exact state he was
+agent = BeeAgent.fromSerialized(json);
+
+// We continue in our conversation
+prompt = "When was he born?";
+console.info(prompt);
+const response2 = await agent.run({
+ prompt,
+});
+console.info(response2.result.text);
diff --git a/examples/agents/simple.ts b/examples/agents/simple.ts
new file mode 100644
index 00000000..9566bcbb
--- /dev/null
+++ b/examples/agents/simple.ts
@@ -0,0 +1,23 @@
+import "dotenv/config.js";
+import { BeeAgent } from "@/agents/bee/agent.js";
+import { TokenMemory } from "@/memory/tokenMemory.js";
+import { DuckDuckGoSearchTool } from "@/tools/search/duckDuckGoSearch.js";
+import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
+import { OpenMeteoTool } from "@/tools/weather/openMeteo.js";
+
+const llm = new OllamaChatLLM();
+const agent = new BeeAgent({
+ llm,
+ memory: new TokenMemory({ llm }),
+ tools: [new DuckDuckGoSearchTool(), new OpenMeteoTool()],
+});
+
+const response = await agent
+ .run({ prompt: "What's the current weather in Las Vegas?" })
+ .observe((emitter) => {
+ emitter.on("update", async ({ data, update, meta }) => {
+ console.log(`Agent (${update.key}) π€ : `, update.value);
+ });
+ });
+
+console.log(`Agent π€ : `, response.result.text);
diff --git a/examples/helpers/io.ts b/examples/helpers/io.ts
new file mode 100644
index 00000000..bc26dd75
--- /dev/null
+++ b/examples/helpers/io.ts
@@ -0,0 +1,73 @@
+import readline from "node:readline/promises";
+import { stdin, stdout } from "node:process";
+import picocolors from "picocolors";
+import * as R from "remeda";
+import stripAnsi from "strip-ansi";
+
+interface ReadFromConsoleInput {
+ fallback?: string;
+ input?: string;
+ allowEmpty?: boolean;
+}
+
+export function createConsoleReader({
+ fallback,
+ input = "User π€ : ",
+ allowEmpty = false,
+}: ReadFromConsoleInput = {}) {
+ const rl = readline.createInterface({ input: stdin, output: stdout, terminal: true });
+ let isActive = true;
+
+ return {
+ write(role: string, data: string) {
+ rl.write(
+ [role && R.piped(picocolors.red, picocolors.bold)(role), stripAnsi(data ?? "")]
+ .filter(Boolean)
+ .join(" ")
+ .concat("\n"),
+ );
+ },
+ async prompt(): Promise {
+ for await (const { prompt } of this) {
+ return prompt;
+ }
+ process.exit(0);
+ },
+ async *[Symbol.asyncIterator]() {
+ if (!isActive) {
+ return;
+ }
+
+ try {
+ rl.write(
+ `${picocolors.dim(`Interactive session has started. To escape, input 'q' and submit.\n`)}`,
+ );
+
+ for (let iteration = 1, prompt = ""; isActive; iteration++) {
+ prompt = await rl.question(R.piped(picocolors.cyan, picocolors.bold)(input));
+ prompt = stripAnsi(prompt);
+
+ if (prompt === "q") {
+ break;
+ }
+ if (!prompt.trim() || prompt === "\n") {
+ prompt = fallback ?? "";
+ }
+ if (allowEmpty !== false && !prompt.trim()) {
+ rl.write("Error: Empty prompt is not allowed. Please try again.\n");
+ iteration -= 1;
+ continue;
+ }
+ yield { prompt, iteration };
+ }
+ } catch (e) {
+ if (e.code === "ERR_USE_AFTER_CLOSE") {
+ return;
+ }
+ } finally {
+ isActive = false;
+ rl.close();
+ }
+ },
+ };
+}
diff --git a/examples/helpers/setup.ts b/examples/helpers/setup.ts
new file mode 100644
index 00000000..3286e2c5
--- /dev/null
+++ b/examples/helpers/setup.ts
@@ -0,0 +1,4 @@
+import "dotenv/config";
+
+import { Logger } from "@/logger/logger.js";
+Logger.defaults.pretty = true;
diff --git a/examples/llms/chat.ts b/examples/llms/chat.ts
new file mode 100644
index 00000000..79d09cab
--- /dev/null
+++ b/examples/llms/chat.ts
@@ -0,0 +1,22 @@
+import "dotenv/config.js";
+import { createConsoleReader } from "examples/helpers/io.js";
+import { Logger } from "@/logger/logger.js";
+import { BaseMessage, Role } from "@/llms/primitives/message.js";
+import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
+
+Logger.root.level = "info"; // or your custom level
+
+const llm = new OllamaChatLLM();
+
+const reader = createConsoleReader();
+
+for await (const { prompt } of reader) {
+ const response = await llm.generate([
+ BaseMessage.of({
+ role: Role.USER,
+ text: prompt,
+ }),
+ ]);
+ reader.write(`LLM π€ (txt) : `, response.getTextContent());
+ reader.write(`LLM π€ (raw) : `, JSON.stringify(response.finalResult));
+}
diff --git a/examples/llms/chatCallback.ts b/examples/llms/chatCallback.ts
new file mode 100644
index 00000000..8ccc9b90
--- /dev/null
+++ b/examples/llms/chatCallback.ts
@@ -0,0 +1,35 @@
+import "dotenv/config.js";
+import { createConsoleReader } from "examples/helpers/io.js";
+import { Logger } from "@/logger/logger.js";
+import { BaseMessage, Role } from "@/llms/primitives/message.js";
+import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
+
+Logger.root.level = "info"; // or your custom level
+
+const llm = new OllamaChatLLM();
+
+const reader = createConsoleReader();
+
+for await (const { prompt } of reader) {
+ const response = await llm
+ .generate(
+ [
+ BaseMessage.of({
+ role: Role.USER,
+ text: prompt,
+ }),
+ ],
+ {},
+ )
+ .observe((emitter) =>
+ emitter.match("*", (data, event) => {
+ reader.write(`LLM π€ (event: ${event.name})`, JSON.stringify(data));
+
+ // if you want to premature close the stream, just uncomment the following line
+ // callbacks.abort()
+ }),
+ );
+
+ reader.write(`LLM π€ (txt) : `, response.getTextContent());
+ reader.write(`LLM π€ (raw) : `, JSON.stringify(response.finalResult));
+}
diff --git a/examples/llms/chatStream.ts b/examples/llms/chatStream.ts
new file mode 100644
index 00000000..8d779486
--- /dev/null
+++ b/examples/llms/chatStream.ts
@@ -0,0 +1,23 @@
+import "dotenv/config.js";
+import { createConsoleReader } from "examples/helpers/io.js";
+import { Logger } from "@/logger/logger.js";
+import { BaseMessage, Role } from "@/llms/primitives/message.js";
+import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
+
+Logger.root.level = "info"; // or your custom level
+
+const llm = new OllamaChatLLM();
+
+const reader = createConsoleReader();
+
+for await (const { prompt } of reader) {
+ for await (const chunk of llm.stream([
+ BaseMessage.of({
+ role: Role.USER,
+ text: prompt,
+ }),
+ ])) {
+ reader.write(`LLM π€ (txt) : `, chunk.getTextContent());
+ reader.write(`LLM π€ (raw) : `, JSON.stringify(chunk.finalResult));
+ }
+}
diff --git a/examples/llms/providers/bam.ts b/examples/llms/providers/bam.ts
new file mode 100644
index 00000000..5e87dad4
--- /dev/null
+++ b/examples/llms/providers/bam.ts
@@ -0,0 +1,32 @@
+import { BaseMessage } from "@/llms/primitives/message.js";
+import { BAMLLM } from "@/adapters/bam/llm.js";
+import { BAMChatLLM } from "@/adapters/bam/chat.js";
+
+{
+ console.info("===RAW===");
+ const llm = new BAMLLM({
+ modelId: "google/flan-ul2",
+ });
+
+ console.info("Meta", await llm.meta());
+
+ const response = await llm.generate("Hello world!", {
+ stream: true,
+ });
+ console.info(response.finalResult);
+}
+
+{
+ console.info("===CHAT===");
+ const llm = BAMChatLLM.fromPreset("meta-llama/llama-3-1-70b-instruct");
+
+ console.info("Meta", await llm.meta());
+
+ const response = await llm.generate([
+ BaseMessage.of({
+ role: "user",
+ text: "Hello world!",
+ }),
+ ]);
+ console.info(response.messages);
+}
diff --git a/examples/llms/providers/langchain.ts b/examples/llms/providers/langchain.ts
new file mode 100644
index 00000000..3bf48067
--- /dev/null
+++ b/examples/llms/providers/langchain.ts
@@ -0,0 +1,26 @@
+// NOTE: ensure you have installed following packages
+// - @langchain/core
+// - @langchain/cohere (or any other provider related package that you would like to use)
+// List of available providers: https://js.langchain.com/v0.2/docs/integrations/chat/
+
+import { BaseMessage } from "@/llms/primitives/message.js";
+import { LangChainChatLLM } from "@/adapters/langchain/llms/chat.js";
+// @ts-expect-error package not installed
+import { ChatCohere } from "@langchain/cohere";
+
+console.info("===CHAT===");
+const llm = new LangChainChatLLM(
+ new ChatCohere({
+ model: "command-r-plus",
+ temperature: 0,
+ }),
+);
+
+const response = await llm.generate([
+ BaseMessage.of({
+ role: "user",
+ text: "Hello world!",
+ }),
+]);
+console.info(response.messages);
+console.info(response.getTextContent());
diff --git a/examples/llms/providers/ollama.ts b/examples/llms/providers/ollama.ts
new file mode 100644
index 00000000..200b1755
--- /dev/null
+++ b/examples/llms/providers/ollama.ts
@@ -0,0 +1,42 @@
+import { OllamaLLM } from "@/adapters/ollama/llm.js";
+import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
+import { BaseMessage } from "@/llms/primitives/message.js";
+
+{
+ console.info("===RAW===");
+ const llm = new OllamaLLM({
+ modelId: "llama3.1",
+ parameters: {
+ num_predict: 10,
+ stop: ["post"],
+ },
+ });
+
+ console.info("Meta", await llm.meta());
+
+ const response = await llm.generate("Hello world!", {
+ stream: true,
+ });
+ console.info(response.finalResult);
+}
+
+{
+ console.info("===CHAT===");
+ const llm = new OllamaChatLLM({
+ modelId: "llama3.1",
+ parameters: {
+ num_predict: 10,
+ temperature: 0,
+ },
+ });
+
+ console.info("Meta", await llm.meta());
+
+ const response = await llm.generate([
+ BaseMessage.of({
+ role: "user",
+ text: "Hello world!",
+ }),
+ ]);
+ console.info(response.finalResult);
+}
diff --git a/examples/llms/providers/openai.ts b/examples/llms/providers/openai.ts
new file mode 100644
index 00000000..5b28ecbf
--- /dev/null
+++ b/examples/llms/providers/openai.ts
@@ -0,0 +1,20 @@
+import "dotenv/config";
+import { BaseMessage } from "@/llms/primitives/message.js";
+import { OpenAIChatLLM } from "@/adapters/openai/chat.js";
+
+const llm = new OpenAIChatLLM({
+ modelId: "gpt-4o",
+ parameters: {
+ max_tokens: 10,
+ stop: ["post"],
+ },
+});
+
+console.info("Meta", await llm.meta());
+const response = await llm.generate([
+ BaseMessage.of({
+ role: "user",
+ text: "Hello world!",
+ }),
+]);
+console.info(response.getTextContent());
diff --git a/examples/llms/providers/watsonx.ts b/examples/llms/providers/watsonx.ts
new file mode 100644
index 00000000..6971ada9
--- /dev/null
+++ b/examples/llms/providers/watsonx.ts
@@ -0,0 +1,53 @@
+import "dotenv/config";
+import { BaseMessage } from "@/llms/primitives/message.js";
+import { WatsonXChatLLM } from "@/adapters/watsonx/chat.js";
+import { WatsonXLLM } from "@/adapters/watsonx/llm.js";
+import { PromptTemplate } from "@/template.js";
+
+const template = new PromptTemplate({
+ variables: ["messages"],
+ template: `{{#messages}}{{#system}}<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+{{system}}<|eot_id|>{{/system}}{{#user}}<|start_header_id|>user<|end_header_id|>
+
+{{user}}<|eot_id|>{{/user}}{{#assistant}}<|start_header_id|>assistant<|end_header_id|>
+
+{{assistant}}<|eot_id|>{{/assistant}}{{/messages}}<|start_header_id|>assistant<|end_header_id|>
+
+`,
+});
+
+const llm = new WatsonXLLM({
+ modelId: "meta-llama/llama-3-70b-instruct",
+ projectId: process.env.WATSONX_PROJECT_ID,
+ apiKey: process.env.WATSONX_API_KEY,
+ parameters: {
+ decoding_method: "greedy",
+ max_new_tokens: 50,
+ },
+});
+
+const chatLLM = new WatsonXChatLLM({
+ llm,
+ config: {
+ messagesToPrompt(messages: BaseMessage[]) {
+ return template.render({
+ messages: messages.map((message) => ({
+ system: message.role === "system" ? [message.text] : [],
+ user: message.role === "user" ? [message.text] : [],
+ assistant: message.role === "assistant" ? [message.text] : [],
+ })),
+ });
+ },
+ },
+});
+
+console.info("Meta", await chatLLM.meta());
+
+const response = await chatLLM.generate([
+ BaseMessage.of({
+ role: "user",
+ text: "Hello world!",
+ }),
+]);
+console.info(response.messages[0]);
diff --git a/examples/llms/structured.ts b/examples/llms/structured.ts
new file mode 100644
index 00000000..765010ee
--- /dev/null
+++ b/examples/llms/structured.ts
@@ -0,0 +1,28 @@
+import "dotenv/config.js";
+import { z } from "zod";
+import { BaseMessage, Role } from "@/llms/primitives/message.js";
+import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
+
+const llm = new OllamaChatLLM();
+const response = await llm.generateStructured(
+ z.union([
+ z.object({
+ firstName: z.string().min(1),
+ lastName: z.string().min(1),
+ address: z.string(),
+ age: z.number().int().min(1),
+ hobby: z.string(),
+ }),
+ z.object({
+ error: z.string(),
+ }),
+ ]),
+ [
+ BaseMessage.of({
+ role: Role.USER,
+ text: "Generate a profile of a citizen of Europe.", // feel free to update it
+ }),
+ ],
+);
+console.info(response);
+process.exit(0);
diff --git a/examples/llms/text.ts b/examples/llms/text.ts
new file mode 100644
index 00000000..ccc33a13
--- /dev/null
+++ b/examples/llms/text.ts
@@ -0,0 +1,20 @@
+import "dotenv/config.js";
+import { createConsoleReader } from "examples/helpers/io.js";
+import { WatsonXLLM } from "@/adapters/watsonx/llm.js";
+
+const llm = new WatsonXLLM({
+ modelId: "google/flan-ul2",
+ projectId: process.env.WATSONX_PROJECT_ID,
+ apiKey: process.env.WATSONX_API_KEY,
+ parameters: {
+ decoding_method: "greedy",
+ max_new_tokens: 50,
+ },
+});
+
+const reader = createConsoleReader();
+
+const prompt = await reader.prompt();
+const response = await llm.generate(prompt);
+reader.write(`LLM π€ (text) : `, response.getTextContent());
+process.exit(0);
diff --git a/examples/template.ts b/examples/template.ts
new file mode 100644
index 00000000..cc77bc7c
--- /dev/null
+++ b/examples/template.ts
@@ -0,0 +1,68 @@
+import "examples/helpers/setup.js";
+import { PromptTemplate } from "@/template.js";
+import { Logger } from "@/logger/logger.js";
+
+const logger = new Logger({ name: "template" });
+
+// Primitives
+{
+ const greetTemplate = new PromptTemplate({
+ template: `Hello {{name}}`,
+ variables: ["name"],
+ });
+
+ const output = greetTemplate.render({
+ name: "Alex",
+ });
+ logger.info(output); // "Hello Alex!"
+}
+
+// Arrays
+{
+ const template = new PromptTemplate({
+ variables: ["colors"],
+ template: `My Favorite Colors: {{#colors}}{{.}} {{/colors}}`,
+ });
+ const output = template.render({
+ colors: ["Green", "Yellow"],
+ });
+ logger.info(output);
+}
+
+// Objects
+{
+ const template = new PromptTemplate({
+ template: `Expected Duration: {{expected}}ms; Retrieved: {{#responses}}{{duration}}ms {{/responses}}`,
+ variables: ["expected", "responses"],
+ defaults: {
+ expected: 5,
+ },
+ });
+ const output = template.render({
+ expected: PromptTemplate.defaultPlaceholder,
+ responses: [{ duration: 3 }, { duration: 5 }, { duration: 6 }],
+ });
+ logger.info(output);
+}
+
+// Forking
+{
+ const original = new PromptTemplate({
+ template: `You are a helpful assistant called {{name}}. You objective is to {{objective}}.`,
+ variables: ["name", "objective"],
+ });
+
+ const modified = original.fork((oldConfig) => ({
+ ...oldConfig,
+ template: `${oldConfig.template} Your answers must be concise.`,
+ defaults: {
+ name: "Allan",
+ },
+ }));
+
+ const output = modified.render({
+ name: PromptTemplate.defaultPlaceholder,
+ objective: "fulfill the user needs",
+ });
+ logger.info(output);
+}
diff --git a/package.json b/package.json
new file mode 100644
index 00000000..f975f70e
--- /dev/null
+++ b/package.json
@@ -0,0 +1,153 @@
+{
+ "name": "bee-agent-framework",
+ "version": "0.0.0",
+ "license": "Apache-2.0",
+ "description": "Bee - LLM Agent Framework",
+ "author": "IBM Corp.",
+ "contributors": [
+ "Tomas Dvorak "
+ ],
+ "keywords": [
+ "LLM Agent Framework",
+ "Bee Agent Framework",
+ "NodeJS Agent Framework"
+ ],
+ "packageManager": "yarn@4.1.1",
+ "type": "module",
+ "main": "./dist/index.js",
+ "types": "./dist/index.d.ts",
+ "sideEffects": false,
+ "exports": {
+ "./package.json": "./package.json",
+ ".": {
+ "import": {
+ "types": "./dist/index.d.ts",
+ "default": "./dist/index.js"
+ },
+ "require": {
+ "types": "./dist/index.d.cts",
+ "default": "./dist/index.cjs"
+ }
+ },
+ "./*": {
+ "import": {
+ "types": "./dist/*.d.ts",
+ "default": "./dist/*.js"
+ },
+ "require": {
+ "types": "./dist/*.d.cts",
+ "default": "./dist/*.cjs"
+ }
+ }
+ },
+ "files": [
+ "dist/**/*"
+ ],
+ "scripts": {
+ "clean": "rimraf dist",
+ "build": "yarn clean && yarn ts:check && tsup",
+ "ts:check": "tsc --noEmit",
+ "start:bee": "tsx examples/agents/bee.ts",
+ "infra:start-all": "yarn _docker compose up -d",
+ "infra:start-code-interpreter": "yarn _docker compose up bee-code-interpreter",
+ "infra:stop-all": "yarn _docker compose down",
+ "infra:clean-all": "yarn _docker compose down --volumes",
+ "lint": "yarn eslint src examples",
+ "lint:fix": "yarn eslint --fix src examples",
+ "format": "yarn prettier --check src examples",
+ "format:fix": "yarn prettier --write src examples",
+ "test:unit": "vitest run src",
+ "test:unit:watch": "vitest run src",
+ "test:e2e": "vitest run tests",
+ "test:e2e:watch": "vitest watch tests",
+ "test:all": "vitest run",
+ "test:watch": "vitest watch",
+ "prepare": "husky",
+ "copyright": "./scripts/copyright.sh",
+ "release": "release-it",
+ "_ensure_env": "cp -n .env.template .env",
+ "_docker": "yarn _ensure_env && sh -c 'source ./.env && docker_cmd=$(which docker >/dev/null 2>&1 && printf docker || printf podman) && $docker_cmd \"$@\"' sh"
+ },
+ "dependencies": {
+ "@ai-zen/node-fetch-event-source": "^2.1.4",
+ "@connectrpc/connect": "^1.4.0",
+ "@connectrpc/connect-node": "^1.4.0",
+ "ajv": "^8.17.1",
+ "ajv-formats": "^3.0.1",
+ "bee-proto": "0.0.1",
+ "dirty-json": "0.9.2",
+ "duck-duck-scrape": "^2.2.5",
+ "fast-xml-parser": "^4.4.1",
+ "header-generator": "^2.1.54",
+ "joplin-turndown-plugin-gfm": "^1.0.12",
+ "mustache": "^4.2.0",
+ "object-hash": "^3.0.0",
+ "p-queue": "^8.0.1",
+ "p-throttle": "^6.2.0",
+ "pino": "^9.3.2",
+ "promise-based-task": "^3.0.2",
+ "remeda": "^2.11.0",
+ "serialize-error": "^11.0.3",
+ "string-comparison": "^1.3.0",
+ "string-strip-html": "^13.4.8",
+ "turndown": "^7.2.0",
+ "wikipedia": "^2.1.2",
+ "zod": "^3.23.8",
+ "zod-to-json-schema": "^3.23.2"
+ },
+ "peerDependencies": {
+ "@ibm-generative-ai/node-sdk": "~3.2.1",
+ "@langchain/community": "~0.2.28",
+ "@langchain/core": "~0.2.27",
+ "@langchain/langgraph": "~0.0.34",
+ "ollama": "^0.5.8",
+ "openai": "^4.56.0",
+ "openai-chat-tokens": "^0.2.8"
+ },
+ "devDependencies": {
+ "@eslint/js": "^9.9.0",
+ "@ibm-generative-ai/node-sdk": "~3.2.1",
+ "@langchain/community": "~0.2.28",
+ "@langchain/core": "~0.2.27",
+ "@langchain/langgraph": "~0.0.34",
+ "@release-it/conventional-changelog": "^8.0.1",
+ "@rollup/plugin-commonjs": "^26.0.1",
+ "@swc/core": "^1.7.14",
+ "@types/eslint-config-prettier": "^6.11.3",
+ "@types/eslint__js": "^8.42.3",
+ "@types/mustache": "^4",
+ "@types/needle": "^3.3.0",
+ "@types/node": "^20.16.1",
+ "@types/object-hash": "^3.0.6",
+ "@types/turndown": "^5.0.5",
+ "dotenv": "^16.4.5",
+ "eslint": "^9.9.0",
+ "eslint-config-prettier": "^9.1.0",
+ "eslint-plugin-unused-imports": "^4.1.3",
+ "glob": "^11.0.0",
+ "husky": "^9.1.5",
+ "langchain": "~0.2.16",
+ "lint-staged": "^15.2.9",
+ "ollama": "^0.5.8",
+ "openai": "^4.56.0",
+ "openai-chat-tokens": "^0.2.8",
+ "openapi-fetch": "^0.11.1",
+ "openapi-typescript": "^7.3.0",
+ "picocolors": "^1.0.1",
+ "pino-pretty": "^11.2.2",
+ "pino-test": "^1.0.1",
+ "prettier": "^3.3.3",
+ "release-it": "^17.6.0",
+ "rimraf": "^6.0.1",
+ "strip-ansi": "^7.1.0",
+ "temp-dir": "^3.0.0",
+ "tsc-files": "^1.1.4",
+ "tsup": "^8.2.4",
+ "tsx": "^4.17.0",
+ "typescript": "^5.5.4",
+ "typescript-eslint": "^8.2.0",
+ "vite-tsconfig-paths": "^5.0.1",
+ "vitest": "^2.0.5",
+ "yaml": "^2.5.0"
+ }
+}
diff --git a/prettier.config.js b/prettier.config.js
new file mode 100644
index 00000000..ae168549
--- /dev/null
+++ b/prettier.config.js
@@ -0,0 +1,8 @@
+// @ts-check
+
+/** @type {import("prettier").Config} */
+const config = {
+ printWidth: 100,
+};
+
+export default config;
diff --git a/scripts/copyright.sh b/scripts/copyright.sh
new file mode 100755
index 00000000..29b56119
--- /dev/null
+++ b/scripts/copyright.sh
@@ -0,0 +1,34 @@
+#!/bin/bash
+# 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.
+set -e
+
+# Path to the package.json file
+PACKAGE_JSON_PATH="./package.json"
+
+# Check if the package.json file exists
+if [[ ! -f "$PACKAGE_JSON_PATH" ]]; then
+ echo "Error: package.json file not found at $PACKAGE_JSON_PATH"
+ exit 1
+fi
+
+# Retrieve the author property using jq
+AUTHOR=$(jq -r '.author' "$PACKAGE_JSON_PATH")
+
+# Check if the author property is not null or empty
+if [[ ! -n "$AUTHOR" ]]; then
+ echo "Error: author property not found in package.json"
+fi
+
+nwa add -l apache -c "$AUTHOR" src dist tests
diff --git a/src/adapters/bam/__snapshots__/llm.test.ts.snap b/src/adapters/bam/__snapshots__/llm.test.ts.snap
new file mode 100644
index 00000000..ec5434a9
--- /dev/null
+++ b/src/adapters/bam/__snapshots__/llm.test.ts.snap
@@ -0,0 +1,833 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`SDK LLM Output > Merging > Merges moderations 1`] = `
+{
+ "generated_text": "Text 0",
+ "generated_token_count": 3,
+ "input_text": "Hello world!",
+ "input_token_count": 4,
+ "moderations": {
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+ },
+ "seed": 1230,
+ "stop_reason": "max_tokens",
+}
+`;
+
+exports[`SDK LLM Output > Merging > Merges moderations 2`] = `
+{
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 777,
+ "start": 777,
+ },
+ "score": 0.75,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 888,
+ "start": 888,
+ },
+ "score": 0.75,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+}
+`;
+
+exports[`SDK LLM Output > Merging > Multiple chunks (2) 1`] = `
+{
+ "generated_text": "Text 0Text 1",
+ "generated_token_count": 6,
+ "input_text": "Hello world!",
+ "input_token_count": 4,
+ "moderations": {
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+ },
+ "seed": 1230,
+ "stop_reason": "max_tokens",
+}
+`;
+
+exports[`SDK LLM Output > Merging > Multiple chunks (2) 2`] = `
+{
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+}
+`;
+
+exports[`SDK LLM Output > Merging > Multiple chunks (5) 1`] = `
+{
+ "generated_text": "Text 0Text 1Text 2Text 3Text 4",
+ "generated_token_count": 15,
+ "input_text": "Hello world!",
+ "input_token_count": 4,
+ "moderations": {
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+ },
+ "seed": 1230,
+ "stop_reason": "max_tokens",
+}
+`;
+
+exports[`SDK LLM Output > Merging > Multiple chunks (5) 2`] = `
+{
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+}
+`;
+
+exports[`SDK LLM Output > Merging > Multiple chunks (6) 1`] = `
+{
+ "generated_text": "Text 0Text 1Text 2Text 3Text 4Text 5",
+ "generated_token_count": 18,
+ "input_text": "Hello world!",
+ "input_token_count": 4,
+ "moderations": {
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+ },
+ "seed": 1230,
+ "stop_reason": "max_tokens",
+}
+`;
+
+exports[`SDK LLM Output > Merging > Multiple chunks (6) 2`] = `
+{
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+}
+`;
+
+exports[`SDK LLM Output > Merging > Multiple chunks (10) 1`] = `
+{
+ "generated_text": "Text 0Text 1Text 2Text 3Text 4Text 5Text 6Text 7Text 8Text 9",
+ "generated_token_count": 30,
+ "input_text": "Hello world!",
+ "input_token_count": 4,
+ "moderations": {
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 6,
+ "start": 6,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 6,
+ "start": 6,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 6,
+ "start": 6,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 6,
+ "start": 6,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 8,
+ "start": 8,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 8,
+ "start": 8,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+ },
+ "seed": 1230,
+ "stop_reason": "max_tokens",
+}
+`;
+
+exports[`SDK LLM Output > Merging > Multiple chunks (10) 2`] = `
+{
+ "hap": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 6,
+ "start": 6,
+ },
+ "score": 0.98,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 6,
+ "start": 6,
+ },
+ "score": 0.33,
+ "success": true,
+ },
+ ],
+ "social_bias": [
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 0,
+ "start": 0,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 2,
+ "start": 2,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 4,
+ "start": 4,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 6,
+ "start": 6,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 6,
+ "start": 6,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 8,
+ "start": 8,
+ },
+ "score": 0.99,
+ "success": false,
+ },
+ {
+ "flagged": true,
+ "position": {
+ "end": 8,
+ "start": 8,
+ },
+ "score": 0.77,
+ "success": true,
+ },
+ ],
+}
+`;
diff --git a/src/adapters/bam/chat.ts b/src/adapters/bam/chat.ts
new file mode 100644
index 00000000..6227cb36
--- /dev/null
+++ b/src/adapters/bam/chat.ts
@@ -0,0 +1,193 @@
+/**
+ * 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 { AsyncStream, GenerateCallbacks, LLMError } from "@/llms/base.js";
+import { isFunction, isObjectType } from "remeda";
+import {
+ BAMLLM,
+ BAMLLMGenerateOptions,
+ BAMLLMParameters,
+ BAMLLMOutput,
+} from "@/adapters/bam/llm.js";
+import { ChatLLM, ChatLLMOutput } from "@/llms/chat.js";
+import { BaseMessage } from "@/llms/primitives/message.js";
+import { PromptTemplate } from "@/template.js";
+import { Cache } from "@/cache/decoratorCache.js";
+import { BAMChatLLMPreset, BAMChatLLMPresetModel } from "@/adapters/bam/chatPreset.js";
+import { Client } from "@ibm-generative-ai/node-sdk";
+import { transformAsyncIterable } from "@/internals/helpers/stream.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { Emitter } from "@/emitter/emitter.js";
+import { GetRunContext } from "@/context.js";
+
+export class BAMChatLLMOutput extends ChatLLMOutput {
+ public readonly raw: BAMLLMOutput;
+
+ constructor(rawOutput: BAMLLMOutput) {
+ super();
+ this.raw = rawOutput;
+ }
+
+ get finalResult() {
+ return this.raw.finalResult;
+ }
+
+ @Cache()
+ get messages(): BaseMessage[] {
+ const text = this.raw.getTextContent();
+ return [
+ BaseMessage.of({
+ role: "assistant",
+ text,
+ meta: this.raw.meta,
+ }),
+ ];
+ }
+
+ merge(other: BAMChatLLMOutput): void {
+ Cache.getInstance(this, "messages").clear();
+ this.raw.merge(other.raw);
+ }
+
+ getTextContent(): string {
+ const [message] = this.messages;
+ return message.text;
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+
+ createSnapshot() {
+ return {
+ raw: shallowCopy(this.raw),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
+
+export interface BAMChatLLMInputConfig {
+ messagesToPrompt: PromptTemplate<"messages"> | ((messages: BaseMessage[]) => string);
+}
+
+export interface BAMChatLLMInput {
+ llm: BAMLLM;
+ config: BAMChatLLMInputConfig;
+}
+
+export class BAMChatLLM extends ChatLLM {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["bam", "chat_llm"],
+ creator: this,
+ });
+
+ public readonly llm: BAMLLM;
+ protected readonly config: BAMChatLLMInputConfig;
+
+ constructor({ llm, config }: BAMChatLLMInput) {
+ super(llm.modelId, llm.executionOptions);
+ this.llm = llm;
+ this.config = config;
+ }
+
+ static {
+ this.register();
+ }
+
+ async meta() {
+ return this.llm.meta();
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ modelId: this.modelId,
+ executionOptions: shallowCopy(this.executionOptions),
+ llm: this.llm,
+ config: shallowCopy(this.config),
+ };
+ }
+
+ async tokenize(messages: BaseMessage[]) {
+ const prompt = this.messagesToPrompt(messages);
+ return this.llm.tokenize(prompt);
+ }
+
+ protected async _generate(
+ messages: BaseMessage[],
+ options: BAMLLMGenerateOptions,
+ run: GetRunContext,
+ ): Promise {
+ const prompt = this.messagesToPrompt(messages);
+ // @ts-expect-error protected property
+ const rawResponse = await this.llm._generate(prompt, options, run);
+ return new BAMChatLLMOutput(rawResponse);
+ }
+
+ protected async *_stream(
+ messages: BaseMessage[],
+ options: BAMLLMGenerateOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ const prompt = this.messagesToPrompt(messages);
+ // @ts-expect-error protected property
+ const response = this.llm._stream(prompt, options, run);
+ return yield* transformAsyncIterable(response, (output) => new BAMChatLLMOutput(output));
+ }
+
+ messagesToPrompt(messages: BaseMessage[]) {
+ const convertor = this.config.messagesToPrompt;
+ if (convertor instanceof PromptTemplate) {
+ return convertor.render({ messages });
+ }
+ return convertor(messages);
+ }
+
+ static fromPreset(
+ modelId: BAMChatLLMPresetModel,
+ overrides?: {
+ client?: Client;
+ parameters?: BAMLLMParameters | ((value: BAMLLMParameters) => BAMLLMParameters);
+ },
+ ) {
+ const presetFactory = BAMChatLLMPreset[modelId];
+ if (!presetFactory) {
+ throw new LLMError(`Model "${modelId}" does not exist in preset.`);
+ }
+
+ const preset = presetFactory();
+ let parameters = preset.base.parameters ?? {};
+ if (isFunction(overrides?.parameters)) {
+ parameters = overrides?.parameters(parameters);
+ } else if (isObjectType(overrides?.parameters)) {
+ parameters = overrides?.parameters;
+ }
+
+ return new BAMChatLLM({
+ config: preset.chat,
+ llm: new BAMLLM({
+ ...preset.base,
+ ...overrides,
+ parameters,
+ client: overrides?.client ?? new Client(),
+ modelId,
+ }),
+ });
+ }
+}
diff --git a/src/adapters/bam/chatPreset.ts b/src/adapters/bam/chatPreset.ts
new file mode 100644
index 00000000..cf97448e
--- /dev/null
+++ b/src/adapters/bam/chatPreset.ts
@@ -0,0 +1,109 @@
+/**
+ * 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 { BAMChatLLMInputConfig } from "@/adapters/bam/chat.js";
+import { BaseMessage } from "@/llms/primitives/message.js";
+import { PromptTemplate } from "@/template.js";
+import { BAMLLMInput } from "@/adapters/bam/llm.js";
+import { toBoundedFunction } from "@/serializer/utils.js";
+
+interface BAMChatLLMPreset {
+ chat: BAMChatLLMInputConfig;
+ base: Omit;
+}
+
+export const BAMChatLLMPreset = {
+ "meta-llama/llama-3-1-70b-instruct": (): BAMChatLLMPreset => {
+ return {
+ base: {
+ parameters: {
+ decoding_method: "greedy",
+ include_stop_sequence: false,
+ max_new_tokens: 2048,
+ repetition_penalty: 1.03,
+ stop_sequences: ["<|eot_id|>"],
+ },
+ },
+ chat: {
+ messagesToPrompt: toBoundedFunction(
+ (messages: BaseMessage[]) => {
+ const template = new PromptTemplate({
+ variables: ["messages"],
+ template: `{{#messages}}{{#system}}<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+{{system}}<|eot_id|>{{/system}}{{#user}}<|start_header_id|>user<|end_header_id|>
+
+{{user}}<|eot_id|>{{/user}}{{#assistant}}<|start_header_id|>assistant<|end_header_id|>
+
+{{assistant}}<|eot_id|>{{/assistant}}{{#ipython}}<|start_header_id|>ipython<|end_header_id|>
+
+{{ipython}}<|eot_id|>{{/ipython}}{{/messages}}<|start_header_id|>assistant<|end_header_id|>
+`,
+ });
+ return template.render({
+ messages: messages.map((message) => ({
+ system: message.role === "system" ? [message.text] : [],
+ user: message.role === "user" ? [message.text] : [],
+ assistant: message.role === "assistant" ? [message.text] : [],
+ ipython: message.role === "ipython" ? [message.text] : [],
+ })),
+ });
+ },
+ [PromptTemplate],
+ ),
+ },
+ };
+ },
+ "qwen/qwen2-72b-instruct": (): BAMChatLLMPreset => {
+ return {
+ base: {
+ parameters: {
+ decoding_method: "greedy",
+ include_stop_sequence: false,
+ stop_sequences: ["<|im_end|>"],
+ },
+ },
+ chat: {
+ messagesToPrompt: toBoundedFunction(
+ (messages: BaseMessage[]) => {
+ const template = new PromptTemplate({
+ variables: ["messages"],
+ template: `{{#messages}}{{#system}}<|im_start|>system
+{{system}}<|im_end|>
+{{ end }}{{/system}}{{#user}}<|im_start|>user
+{{user}}<|im_end|>
+{{ end }}{{/user}}{{#assistant}}<|im_start|>assistant
+{{assistant}}<|im_end|>
+{{ end }}{{/assistant}}{{/messages}}<|im_start|>assistant
+`,
+ });
+
+ return template.render({
+ messages: messages.map((message) => ({
+ system: message.role === "system" ? [message.text] : [],
+ user: message.role === "user" ? [message.text] : [],
+ assistant: message.role === "assistant" ? [message.text] : [],
+ })),
+ });
+ },
+ [PromptTemplate],
+ ),
+ },
+ };
+ },
+} as const;
+
+export type BAMChatLLMPresetModel = keyof typeof BAMChatLLMPreset;
diff --git a/src/adapters/bam/llm.test.ts b/src/adapters/bam/llm.test.ts
new file mode 100644
index 00000000..53f13f90
--- /dev/null
+++ b/src/adapters/bam/llm.test.ts
@@ -0,0 +1,163 @@
+/**
+ * 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 { BAMLLMOutput, BAMLLMOutputModeration, BAMLLMOutputResult } from "@/adapters/bam/llm.js";
+import { expect } from "vitest";
+import { LLMOutputError } from "@/llms/base.js";
+import { verifyDeserialization } from "@tests/e2e/utils.js";
+
+describe("SDK LLM Output", () => {
+ const generateModeration = (
+ i: number,
+ override: Record = {},
+ ): NonNullable["hap"]>[number] => {
+ return {
+ success: true,
+ position: { start: i, end: i },
+ score: 0.75,
+ flagged: true,
+ ...override,
+ };
+ };
+
+ const generateChunks = (count: number): BAMLLMOutputResult[] => {
+ return Array(count)
+ .fill(null)
+ .map(
+ (_, i): BAMLLMOutputResult => ({
+ generated_text: `Text ${i}`,
+ generated_token_count: 3,
+ moderations: undefined,
+ ...(i % 2 === 0 && {
+ moderations: {
+ ...(i % 3 === 0 && {
+ hap: [
+ generateModeration(i, { score: 0.98, success: false }),
+ generateModeration(i, { score: 0.33 }),
+ ],
+ }),
+ social_bias: [
+ generateModeration(i, { score: 0.99, success: false }),
+ generateModeration(i, { score: 0.77 }),
+ ],
+ },
+ }),
+ stop_reason: "not_finished",
+ ...(i === 0 && {
+ input_text: "Hello world!",
+ input_token_count: 4,
+ }),
+ ...(i + 1 === count && {
+ stop_reason: "max_tokens",
+ }),
+ seed: 1230,
+ }),
+ );
+ };
+
+ const getInstance = (
+ chunks: BAMLLMOutputResult[],
+ moderations?: BAMLLMOutputModeration | BAMLLMOutputModeration[],
+ ) => {
+ return new BAMLLMOutput({
+ results: chunks,
+ moderations,
+ meta: {
+ id: "X",
+ model_id: "model",
+ input_parameters: undefined,
+ created_at: new Date().toISOString(),
+ },
+ });
+ };
+
+ describe("Merging", () => {
+ it("Throws for no chunks", () => {
+ const chunks = generateChunks(0);
+ const instance = getInstance(chunks);
+ expect(() => instance.finalResult).toThrowError(LLMOutputError);
+ });
+
+ it("Single instance", () => {
+ const chunks = generateChunks(1);
+ const instance = getInstance(chunks);
+ expect(instance.finalResult).toStrictEqual(chunks[0]);
+ });
+
+ it.each([2, 5, 6, 10])("Multiple chunks (%i)", (chunksCount) => {
+ const chunks = generateChunks(chunksCount);
+ const instance = getInstance(chunks);
+
+ const finalResult = instance.finalResult;
+ expect(finalResult).toMatchSnapshot();
+
+ const finalModerations = instance.finalModeration;
+ expect(finalModerations).toMatchSnapshot();
+ });
+
+ it("Merges moderations", () => {
+ const chunks = generateChunks(1);
+ const moderations = {
+ hap: [generateModeration(777)],
+ social_bias: [generateModeration(888)],
+ };
+ const instance = getInstance(chunks, moderations);
+
+ const finalResult = instance.finalResult;
+ expect(finalResult).toMatchSnapshot();
+
+ const finalModerations = instance.finalModeration;
+ expect(finalModerations).toMatchSnapshot();
+ });
+ });
+
+ describe("Caching", () => {
+ it("Caches final result", () => {
+ const [firstChunk, ...chunks] = generateChunks(5);
+ const instance = getInstance([firstChunk, ...chunks]);
+
+ const result = instance.finalResult;
+ expect(result).toBeTruthy();
+ Object.defineProperty(firstChunk, "generated_text", {
+ get() {
+ throw new Error("This should not be called!");
+ },
+ });
+ expect(instance.finalResult).toBe(result);
+ expect(instance.getTextContent()).toBeTruthy();
+ });
+
+ it("Revalidates cache", () => {
+ const instance = getInstance(generateChunks(5));
+
+ const result = instance.finalResult;
+ expect(instance.finalResult).toBe(result);
+ expect(instance.finalResult).toBe(result);
+ instance.merge(getInstance(generateChunks(1)));
+
+ const newResult = instance.finalResult;
+ expect(newResult).not.toBe(result);
+ expect(instance.finalResult).toBe(newResult);
+ });
+ });
+
+ it("Serializes", () => {
+ const instance = getInstance(generateChunks(5));
+ const serialized = instance.serialize();
+ const deserialized = BAMLLMOutput.fromSerialized(serialized);
+ verifyDeserialization(instance, deserialized);
+ });
+});
diff --git a/src/adapters/bam/llm.ts b/src/adapters/bam/llm.ts
new file mode 100644
index 00000000..42da6224
--- /dev/null
+++ b/src/adapters/bam/llm.ts
@@ -0,0 +1,369 @@
+/**
+ * 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 { LLM, LLMInput } from "@/llms/index.js";
+import {
+ AsyncStream,
+ BaseLLMOutput,
+ BaseLLMTokenizeOutput,
+ ExecutionOptions,
+ GenerateCallbacks,
+ GenerateOptions,
+ LLMError,
+ LLMMeta,
+ LLMOutputError,
+} from "@/llms/base.js";
+import {
+ Client,
+ TextGenerationCreateInput,
+ TextGenerationCreateStreamInput,
+ TextGenerationCreateStreamOutput,
+ TextGenerationCreateOutput,
+ HttpError,
+} from "@ibm-generative-ai/node-sdk";
+import * as R from "remeda";
+import { ExcludeNonStringIndex } from "@/internals/types.js";
+import { FrameworkError, NotImplementedError } from "@/errors.js";
+import { Cache } from "@/cache/decoratorCache.js";
+import { transformAsyncIterable } from "@/internals/helpers/stream.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { safeSum } from "@/internals/helpers/number.js";
+import { customMerge, omitUndefined } from "@/internals/helpers/object.js";
+import { isEmpty, isString } from "remeda";
+import { Emitter } from "@/emitter/emitter.js";
+import { GetRunContext } from "@/context.js";
+
+export type BAMLLMOutputMeta = Omit, "results">;
+
+export type BAMLLMOutputResult = ExcludeNonStringIndex<
+ TextGenerationCreateOutput["results"][number]
+>;
+
+export type BAMLLMOutputModeration = ExcludeNonStringIndex<
+ TextGenerationCreateStreamOutput["moderations"]
+>;
+
+export interface BAMLLMOutputConstructor {
+ meta: BAMLLMOutputMeta;
+ results: BAMLLMOutputResult[];
+ moderations?: BAMLLMOutputModeration | BAMLLMOutputModeration[];
+}
+
+export class BAMLLMOutput extends BaseLLMOutput {
+ public readonly meta: BAMLLMOutputMeta;
+ public readonly results: BAMLLMOutputResult[];
+ public readonly moderations: BAMLLMOutputModeration[];
+
+ constructor(content: BAMLLMOutputConstructor) {
+ super();
+ this.meta = content.meta;
+ this.results = content.results;
+ this.moderations = R.isArray(content.moderations)
+ ? content.moderations
+ : [content.moderations].filter(R.isDefined);
+ }
+
+ static {
+ this.register();
+ }
+
+ getTextContent(): string {
+ return this.finalResult.generated_text;
+ }
+
+ get finalModeration(): Readonly {
+ return BAMLLMOutput._combineModerations(...this.moderations, this.finalResult.moderations);
+ }
+
+ @Cache()
+ get finalResult(): Readonly {
+ if (this.results.length === 0) {
+ throw new LLMOutputError("No chunks to get final result from!");
+ }
+
+ return customMerge(this.results, {
+ generated_text: (value = "", oldValue = "") => oldValue + value,
+ input_token_count: safeSum,
+ generated_token_count: safeSum,
+ input_text: (value, oldValue) => value ?? oldValue,
+ generated_tokens: (value, oldValue) => [...(value || []), ...(oldValue || [])],
+ seed: (value, oldValue) => value ?? oldValue,
+ stop_reason: (value, oldValue) => value ?? oldValue,
+ stop_sequence: (value, oldValue) => value ?? oldValue,
+ input_tokens: (value, oldValue) => value ?? oldValue,
+ moderations: (value, oldValue) =>
+ value && oldValue ? BAMLLMOutput._combineModerations(oldValue, value) : (value ?? oldValue),
+ });
+ }
+
+ merge(other: BAMLLMOutput): void {
+ Cache.getInstance(this, "finalResult").clear();
+
+ this.results.push(...other.results);
+ this.moderations.push(...other.moderations);
+ Object.assign(this.meta, omitUndefined(other.meta));
+ }
+
+ createSnapshot() {
+ return {
+ results: shallowCopy(this.results),
+ moderations: shallowCopy(this.moderations),
+ meta: shallowCopy(this.meta),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+
+ protected static _combineModerations(...entries: BAMLLMOutputModeration[]) {
+ const newModerations: NonNullable = {};
+ for (const entry of entries) {
+ for (const [key, records] of R.entries(entry ?? {})) {
+ if (R.isEmpty(records)) {
+ continue;
+ }
+ if (!newModerations[key]) {
+ newModerations[key] = [];
+ }
+ newModerations[key]!.push(...records);
+ }
+ }
+ return newModerations;
+ }
+}
+
+export type BAMLLMParameters = NonNullable<
+ TextGenerationCreateInput["parameters"] & TextGenerationCreateStreamInput["parameters"]
+>;
+
+export interface BAMLLMGenerateOptions extends GenerateOptions {
+ moderations?: TextGenerationCreateInput["moderations"];
+}
+
+export interface BAMLLMInput {
+ client?: Client;
+ modelId: string;
+ parameters?: BAMLLMParameters;
+ executionOptions?: ExecutionOptions;
+}
+
+export class BAMLLM extends LLM {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["bam", "llm"],
+ creator: this,
+ });
+
+ public readonly client: Client;
+ public readonly parameters: Partial;
+
+ constructor({ client, parameters, modelId, executionOptions = {} }: BAMLLMInput) {
+ super(modelId, executionOptions);
+ this.client = client ?? new Client();
+ this.parameters = parameters ?? {};
+ }
+
+ static {
+ this.register();
+ }
+
+ async meta(): Promise {
+ try {
+ const { result } = await this.client.model.retrieve({
+ id: this.modelId,
+ });
+
+ const tokenLimit = result.token_limits?.find?.((limit) => {
+ if (this.parameters?.beam_width !== undefined) {
+ return limit.token_limit !== undefined && limit.beam_width === this.parameters.beam_width;
+ }
+ return limit.token_limit !== undefined;
+ });
+ return {
+ tokenLimit: tokenLimit?.token_limit ?? Infinity,
+ };
+ } catch {
+ // TODO: remove once retrieval gets fixed on the API
+ if (
+ this.modelId === "qwen/qwen2-72b-instruct" ||
+ this.modelId === "meta-llama/llama-3-1-70b-instruct"
+ ) {
+ return {
+ tokenLimit: 131_072,
+ };
+ } else if (this.modelId.includes("llama-3-70b-instruct")) {
+ return {
+ tokenLimit: 8196,
+ };
+ }
+
+ return {
+ tokenLimit: Infinity,
+ };
+ }
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ client: null,
+ modelId: this.modelId,
+ parameters: shallowCopy(this.parameters),
+ executionOptions: shallowCopy(this.executionOptions),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType): void {
+ super.loadSnapshot(snapshot);
+ Object.assign(this, snapshot, {
+ client: snapshot?.client ?? new Client(), // TODO: serialize?
+ });
+ }
+
+ protected _transformError(error: Error): Error {
+ if (error instanceof FrameworkError) {
+ throw error;
+ }
+ if (error instanceof HttpError) {
+ throw new LLMError("LLM has occurred an error!", [error], {
+ isRetryable: [408, 425, 429, 500, 503].includes(error.status_code),
+ });
+ }
+ return new LLMError("LLM has occurred an error!", [error]);
+ }
+
+ async tokenize(input: LLMInput): Promise {
+ try {
+ const {
+ results: [result],
+ } = await this.client.text.tokenization.create({
+ input,
+ model_id: this.modelId,
+ parameters: {
+ return_options: {
+ tokens: true,
+ },
+ },
+ });
+
+ return {
+ tokensCount: result.token_count,
+ tokens: result.tokens,
+ };
+ } catch (e) {
+ throw this._transformError(e);
+ }
+ }
+
+ protected async _generate(
+ input: LLMInput,
+ options: BAMLLMGenerateOptions,
+ run: GetRunContext,
+ ): Promise {
+ try {
+ const response = await this.client.text.generation.create(
+ {
+ input,
+ moderations: options?.moderations,
+ model_id: this.modelId,
+ parameters: this._prepareParameters(options),
+ },
+ {
+ signal: run.signal,
+ },
+ );
+ return this._rawResponseToOutput(response);
+ } catch (e) {
+ throw this._transformError(e);
+ }
+ }
+
+ protected async *_stream(
+ input: string,
+ options: BAMLLMGenerateOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ try {
+ const response = await this.client.text.generation.create_stream(
+ {
+ input,
+ moderations: options?.moderations,
+ model_id: this.modelId,
+ parameters: this._prepareParameters(options),
+ },
+ {
+ signal: run.signal,
+ },
+ );
+ yield* transformAsyncIterable(
+ response[Symbol.asyncIterator](),
+ this._rawResponseToOutput.bind(this),
+ );
+ } catch (e) {
+ throw this._transformError(e);
+ }
+ }
+
+ protected _rawResponseToOutput(
+ raw: TextGenerationCreateOutput | TextGenerationCreateStreamOutput,
+ ) {
+ const chunks = (raw.results ?? []) as BAMLLMOutputResult[];
+
+ return new BAMLLMOutput({
+ results: chunks,
+ moderations: (raw as TextGenerationCreateStreamOutput)?.moderations,
+ meta: R.pickBy(
+ {
+ id: raw.id!,
+ model_id: raw.model_id,
+ created_at: raw.created_at!,
+ input_parameters: raw.input_parameters,
+ },
+ R.isDefined,
+ ),
+ });
+ }
+
+ protected _prepareParameters(overrides?: GenerateOptions): typeof this.parameters {
+ const guided = overrides?.guided ? {} : (this.parameters.guided ?? {});
+ const guidedOverride = overrides?.guided;
+
+ if (guidedOverride?.choice) {
+ guided.choice = guidedOverride.choice;
+ } else if (guidedOverride?.grammar) {
+ guided.grammar = guidedOverride.grammar;
+ } else if (guidedOverride?.json) {
+ guided.json_schema = isString(guidedOverride.json)
+ ? JSON.parse(guidedOverride.json)
+ : guidedOverride.json;
+ } else if (guidedOverride?.regex) {
+ guided.regex = guidedOverride.regex;
+ } else if (!isEmpty(guidedOverride ?? {})) {
+ throw new NotImplementedError(
+ `Following types ${Object.keys(overrides!.guided!).join(",")}" for the constraint decoding are not supported!`,
+ );
+ }
+
+ return {
+ ...this.parameters,
+ guided: isEmpty(guided) ? undefined : guided,
+ };
+ }
+}
diff --git a/src/adapters/langchain/llms/chat.ts b/src/adapters/langchain/llms/chat.ts
new file mode 100644
index 00000000..4396fd1f
--- /dev/null
+++ b/src/adapters/langchain/llms/chat.ts
@@ -0,0 +1,225 @@
+/**
+ * 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 {
+ AsyncStream,
+ BaseLLMTokenizeOutput,
+ ExecutionOptions,
+ GenerateCallbacks,
+ GenerateOptions,
+ LLMMeta,
+} from "@/llms/base.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { load } from "@langchain/core/load";
+import {
+ BaseChatModel,
+ BaseChatModelCallOptions,
+} from "@langchain/core/language_models/chat_models";
+import { ChatLLM, ChatLLMOutput } from "@/llms/chat.js";
+import { BaseMessage, Role, RoleType } from "@/llms/primitives/message.js";
+import {
+ BaseMessageChunk,
+ BaseMessage as LCBaseMessage,
+ ChatMessage as LCMChatMessage,
+ MessageContentComplex,
+ MessageContentText,
+ MessageType,
+} from "@langchain/core/messages";
+import { Cache } from "@/cache/decoratorCache.js";
+import { getProp, omitUndefined } from "@/internals/helpers/object.js";
+import { Emitter } from "@/emitter/emitter.js";
+import { GetRunContext } from "@/context.js";
+
+export class LangChainChatLLMOutput extends ChatLLMOutput {
+ constructor(
+ public messages: BaseMessage[],
+ public meta: Record = {},
+ ) {
+ super();
+ }
+
+ static {
+ this.register();
+ }
+
+ merge(other: LangChainChatLLMOutput): void {
+ this.messages.push(...other.messages);
+ Object.assign(this.meta, omitUndefined(other.meta));
+ }
+
+ getTextContent(): string {
+ return this.messages.map((msg) => msg.text).join("");
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+
+ createSnapshot() {
+ return {
+ messages: shallowCopy(this.messages),
+ meta: shallowCopy(this.meta),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType): void {
+ Object.assign(this, snapshot);
+ }
+}
+
+export type LangChainChatLLMParameters = Record;
+type MergedCallOptions = { lc: T } & GenerateOptions;
+
+export class LangChainChatLLM<
+ CallOptions extends BaseChatModelCallOptions,
+ OutputMessageType extends BaseMessageChunk,
+> extends ChatLLM> {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["langchain", "chat_llm"],
+ creator: this,
+ });
+ public readonly parameters: any;
+
+ constructor(
+ public readonly lcLLM: BaseChatModel,
+ protected modelMeta?: LLMMeta,
+ executionOptions?: ExecutionOptions,
+ ) {
+ super(lcLLM._modelType(), executionOptions);
+ this.parameters = lcLLM.invocationParams();
+ }
+
+ static {
+ this.register();
+ }
+
+ async meta() {
+ if (this.modelMeta) {
+ return this.modelMeta;
+ }
+
+ return {
+ tokenLimit: Infinity,
+ };
+ }
+
+ async tokenize(input: BaseMessage[]): Promise {
+ return {
+ tokensCount: await this.lcLLM.getNumTokens(input),
+ };
+ }
+
+ @Cache()
+ protected get mappers() {
+ const roleMapper = new Map([
+ ["system", Role.SYSTEM],
+ ["assistant", Role.ASSISTANT],
+ ["ai", Role.ASSISTANT],
+ ["generic", Role.ASSISTANT],
+ ["function", Role.ASSISTANT],
+ ["tool", Role.ASSISTANT],
+ ["human", Role.USER],
+ ["tool", Role.ASSISTANT],
+ ]);
+
+ return {
+ toLCMessage(message: BaseMessage): LCBaseMessage {
+ return new LCMChatMessage({
+ role: message.role,
+ content: message.text,
+ response_metadata: message.meta,
+ });
+ },
+ fromLCMessage(message: LCBaseMessage | LCMChatMessage): BaseMessage {
+ const role: string = getProp(message, ["role"], message._getType());
+ const text: string =
+ typeof message.content === "string"
+ ? message.content
+ : message.content
+ .filter(
+ (msg: MessageContentComplex): msg is MessageContentText => msg.type === "text",
+ )
+ .map((msg: MessageContentText) => msg.text)
+ .join("\n");
+
+ return BaseMessage.of({
+ role: roleMapper.has(role) ? roleMapper.get(role)! : Role.ASSISTANT,
+ text,
+ });
+ },
+ };
+ }
+
+ protected async _generate(
+ input: BaseMessage[],
+ options: MergedCallOptions,
+ run: GetRunContext,
+ ): Promise {
+ const lcMessages = input.map((msg) => this.mappers.toLCMessage(msg));
+ const response = await this.lcLLM.invoke(lcMessages, {
+ ...options?.lc,
+ signal: run.signal,
+ });
+
+ return new LangChainChatLLMOutput(
+ [this.mappers.fromLCMessage(response)],
+ response.response_metadata,
+ );
+ }
+
+ protected async *_stream(
+ input: BaseMessage[],
+ options: MergedCallOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ const lcMessages = input.map((msg) => this.mappers.toLCMessage(msg));
+ const response = this.lcLLM._streamResponseChunks(lcMessages, {
+ ...options?.lc,
+ signal: run.signal,
+ });
+ for await (const chunk of response) {
+ yield new LangChainChatLLMOutput(
+ [this.mappers.fromLCMessage(chunk.message)],
+ chunk.message.response_metadata,
+ );
+ }
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ modelId: this.modelId,
+ modelMeta: this.modelMeta,
+ parameters: shallowCopy(this.parameters),
+ executionOptions: shallowCopy(this.executionOptions),
+ lcLLM: JSON.stringify(this.lcLLM.toJSON()),
+ };
+ }
+
+ async loadSnapshot({ lcLLM, ...state }: ReturnType) {
+ super.loadSnapshot(state);
+ Object.assign(this, state, {
+ lcLLM: await (async () => {
+ if (lcLLM.includes("@ibm-generative-ai/node-sdk")) {
+ const { GenAIChatModel } = await import("@ibm-generative-ai/node-sdk/langchain");
+ return GenAIChatModel.fromJSON(lcLLM);
+ }
+
+ return await load(lcLLM);
+ })(),
+ });
+ }
+}
diff --git a/src/adapters/langchain/llms/index.ts b/src/adapters/langchain/llms/index.ts
new file mode 100644
index 00000000..6831bc63
--- /dev/null
+++ b/src/adapters/langchain/llms/index.ts
@@ -0,0 +1,17 @@
+/**
+ * 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.
+ */
+
+export * from "./llm.js";
diff --git a/src/adapters/langchain/llms/llm.test.ts b/src/adapters/langchain/llms/llm.test.ts
new file mode 100644
index 00000000..cc2309bb
--- /dev/null
+++ b/src/adapters/langchain/llms/llm.test.ts
@@ -0,0 +1,41 @@
+/**
+ * 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 { verifyDeserialization } from "@tests/e2e/utils.js";
+import { LangChainLLM } from "@/adapters/langchain/llms/llm.js";
+import { GenAIModel } from "@ibm-generative-ai/node-sdk/langchain";
+import { Client } from "@ibm-generative-ai/node-sdk";
+
+describe("Langchain LLM", () => {
+ const getInstance = () => {
+ return new LangChainLLM(
+ new GenAIModel({
+ model_id: "google/flan-ul2",
+ client: new Client(),
+ parameters: {
+ max_new_tokens: 100,
+ },
+ }),
+ );
+ };
+
+ it("Serializes", async () => {
+ const instance = getInstance();
+ const serialized = instance.serialize();
+ const deserialized = await LangChainLLM.fromSerialized(serialized);
+ verifyDeserialization(instance, deserialized);
+ });
+});
diff --git a/src/adapters/langchain/llms/llm.ts b/src/adapters/langchain/llms/llm.ts
new file mode 100644
index 00000000..51cb527b
--- /dev/null
+++ b/src/adapters/langchain/llms/llm.ts
@@ -0,0 +1,156 @@
+/**
+ * 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 { LLM, LLMInput } from "@/llms/index.js";
+import { BaseLLM as LCBaseLLM } from "@langchain/core/language_models/llms";
+import {
+ AsyncStream,
+ BaseLLMOutput,
+ BaseLLMTokenizeOutput,
+ ExecutionOptions,
+ GenerateCallbacks,
+ InternalGenerateOptions,
+ LLMMeta,
+ StreamGenerateOptions,
+} from "@/llms/base.js";
+import { load } from "@langchain/core/load";
+import { assign } from "@/internals/helpers/object.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { Emitter } from "@/emitter/emitter.js";
+import { GetRunContext } from "@/context.js";
+
+export class LangChainLLMOutput extends BaseLLMOutput {
+ constructor(
+ public text: string,
+ public readonly meta: Record,
+ ) {
+ super();
+ }
+
+ static {
+ this.register();
+ }
+
+ merge(other: LangChainLLMOutput): void {
+ this.text += other.text;
+ assign(this.meta, other.meta);
+ }
+
+ getTextContent(): string {
+ return this.text;
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+
+ createSnapshot() {
+ return {
+ text: this.text,
+ meta: shallowCopy(this.meta),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
+
+export class LangChainLLM extends LLM {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["langchain", "llm"],
+ creator: this,
+ });
+ protected readonly parameters: any;
+
+ constructor(
+ public readonly lcLLM: LCBaseLLM,
+ private modelMeta?: LLMMeta,
+ executionOptions?: ExecutionOptions,
+ ) {
+ super(lcLLM._modelType(), executionOptions);
+ this.parameters = lcLLM.invocationParams();
+ }
+
+ static {
+ this.register();
+ }
+
+ async meta() {
+ if (this.modelMeta) {
+ return this.modelMeta;
+ }
+
+ return {
+ tokenLimit: Infinity,
+ };
+ }
+
+ async tokenize(input: LLMInput): Promise {
+ return {
+ tokensCount: await this.lcLLM.getNumTokens(input),
+ };
+ }
+
+ protected async _generate(
+ input: LLMInput,
+ options: InternalGenerateOptions,
+ run: GetRunContext,
+ ): Promise {
+ const { generations } = await this.lcLLM.generate([input], {
+ signal: run.signal,
+ });
+ return new LangChainLLMOutput(generations[0][0].text, generations[0][0].generationInfo || {});
+ }
+
+ protected async *_stream(
+ input: string,
+ options: StreamGenerateOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ const response = this.lcLLM._streamResponseChunks(input, {
+ signal: run.signal,
+ });
+ for await (const chunk of response) {
+ yield new LangChainLLMOutput(chunk.text, chunk.generationInfo || {});
+ }
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ modelId: this.modelId,
+ modelMeta: this.modelMeta,
+ parameters: shallowCopy(this.parameters),
+ executionOptions: shallowCopy(this.executionOptions),
+ lcLLM: JSON.stringify(this.lcLLM.toJSON()),
+ };
+ }
+
+ async loadSnapshot({ lcLLM, ...state }: ReturnType) {
+ super.loadSnapshot(state);
+ Object.assign(this, state, {
+ lcLLM: await (async () => {
+ if (lcLLM.includes("@ibm-generative-ai/node-sdk")) {
+ const { GenAIModel } = await import("@ibm-generative-ai/node-sdk/langchain");
+ return GenAIModel.fromJSON(lcLLM);
+ }
+
+ return await load(lcLLM);
+ })(),
+ });
+ }
+}
diff --git a/src/adapters/ollama/chat.ts b/src/adapters/ollama/chat.ts
new file mode 100644
index 00000000..dda554da
--- /dev/null
+++ b/src/adapters/ollama/chat.ts
@@ -0,0 +1,215 @@
+/**
+ * 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 {
+ AsyncStream,
+ BaseLLMTokenizeOutput,
+ ExecutionOptions,
+ GenerateCallbacks,
+ GenerateOptions,
+ LLMOutputError,
+ StreamGenerateOptions,
+} from "@/llms/base.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { ChatLLM, ChatLLMOutput } from "@/llms/chat.js";
+import { BaseMessage } from "@/llms/primitives/message.js";
+import { Emitter } from "@/emitter/emitter.js";
+import { ChatResponse, Ollama as Client, Options as Parameters } from "ollama";
+import { signalRace } from "@/internals/helpers/promise.js";
+import { GetRunContext } from "@/context.js";
+import { Cache } from "@/cache/decoratorCache.js";
+import { customMerge } from "@/internals/helpers/object.js";
+import { safeSum } from "@/internals/helpers/number.js";
+import { extractModelMeta, registerClient } from "@/adapters/ollama/shared.js";
+
+export class OllamaChatLLMOutput extends ChatLLMOutput {
+ public readonly results: ChatResponse[];
+
+ constructor(response: ChatResponse) {
+ super();
+ this.results = [response];
+ }
+
+ static {
+ this.register();
+ }
+
+ get messages() {
+ return this.results.flatMap((response) =>
+ BaseMessage.of({
+ role: response.message.role,
+ text: response.message.content,
+ }),
+ );
+ }
+
+ getTextContent(): string {
+ return this.finalResult.message.content;
+ }
+
+ @Cache()
+ get finalResult(): Readonly {
+ if (this.results.length === 0) {
+ throw new LLMOutputError("No chunks to get final result from!");
+ }
+
+ return customMerge(this.results, {
+ message: (value, oldValue) => ({
+ role: value.role ?? oldValue.role,
+ content: `${oldValue?.content ?? ""}${value?.content ?? ""}`,
+ images: [...(oldValue?.images ?? []), ...(value?.images ?? [])] as string[],
+ tool_calls: [...(oldValue?.tool_calls ?? []), ...(value?.tool_calls ?? [])],
+ }),
+ total_duration: (value, oldValue) => value ?? oldValue,
+ load_duration: (value, oldValue) => value ?? oldValue,
+ model: (value, oldValue) => value ?? oldValue,
+ done: (value, oldValue) => value ?? oldValue,
+ done_reason: (value, oldValue) => value ?? oldValue,
+ created_at: (value, oldValue) => value ?? oldValue,
+ eval_duration: (value, oldValue) => value ?? oldValue,
+ prompt_eval_duration: (value, oldValue) => value ?? oldValue,
+ prompt_eval_count: safeSum,
+ eval_count: safeSum,
+ });
+ }
+
+ merge(other: OllamaChatLLMOutput): void {
+ Cache.getInstance(this, "finalResult").clear();
+ this.results.push(...other.results);
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+
+ createSnapshot() {
+ return {
+ results: shallowCopy(this.results),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType): void {
+ Object.assign(this, snapshot);
+ }
+}
+
+interface Input {
+ modelId: string;
+ client?: Client;
+ parameters?: Partial;
+ executionOptions?: ExecutionOptions;
+}
+
+export class OllamaChatLLM extends ChatLLM {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["ollama", "chat_llm"],
+ creator: this,
+ });
+
+ public readonly client: Client;
+ public readonly parameters: Partial;
+
+ constructor(
+ { client, modelId, parameters, executionOptions = {} }: Input = {
+ modelId: "llama3.1",
+ parameters: {
+ temperature: 0,
+ },
+ },
+ ) {
+ super(modelId, executionOptions);
+ this.client = client ?? new Client({ fetch });
+ this.parameters = parameters ?? {};
+ }
+
+ static {
+ this.register();
+ registerClient();
+ }
+
+ async meta() {
+ const model = await this.client.show({
+ model: this.modelId,
+ });
+
+ return extractModelMeta(model);
+ }
+
+ async tokenize(input: BaseMessage[]): Promise {
+ const contentLength = input.reduce((acc, msg) => acc + msg.text.length, 0);
+
+ return {
+ tokensCount: Math.ceil(contentLength / 4),
+ };
+ }
+
+ protected async _generate(
+ input: BaseMessage[],
+ options: GenerateOptions,
+ run: GetRunContext,
+ ): Promise {
+ const response = await signalRace(
+ () =>
+ this.client.chat({
+ model: this.modelId,
+ stream: false,
+ messages: input.map((msg) => ({
+ role: msg.role,
+ content: msg.text,
+ })),
+ options: this.parameters,
+ format: options.guided?.json ? "json" : undefined,
+ }),
+ run.signal,
+ () => this.client.abort(),
+ );
+
+ return new OllamaChatLLMOutput(response);
+ }
+
+ protected async *_stream(
+ input: BaseMessage[],
+ options: StreamGenerateOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ for await (const chunk of await this.client.chat({
+ model: this.modelId,
+ stream: true,
+ messages: input.map((msg) => ({
+ role: msg.role,
+ content: msg.text,
+ })),
+ options: this.parameters,
+ format: options.guided?.json ? "json" : undefined,
+ })) {
+ if (run.signal.aborted) {
+ break;
+ }
+ yield new OllamaChatLLMOutput(chunk);
+ }
+ run.signal.throwIfAborted();
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ modelId: this.modelId,
+ parameters: shallowCopy(this.parameters),
+ executionOptions: shallowCopy(this.executionOptions),
+ client: this.client,
+ };
+ }
+}
diff --git a/src/adapters/ollama/llm.test.ts b/src/adapters/ollama/llm.test.ts
new file mode 100644
index 00000000..dbfe1841
--- /dev/null
+++ b/src/adapters/ollama/llm.test.ts
@@ -0,0 +1,49 @@
+/**
+ * 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 { verifyDeserialization } from "@tests/e2e/utils.js";
+import { OllamaLLM } from "@/adapters/ollama/llm.js";
+import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
+
+describe("Ollama LLM", () => {
+ const getInstance = () => {
+ return new OllamaLLM({
+ modelId: "llama3.1",
+ });
+ };
+
+ it("Serializes", async () => {
+ const instance = getInstance();
+ const serialized = instance.serialize();
+ const deserialized = OllamaLLM.fromSerialized(serialized);
+ verifyDeserialization(instance, deserialized);
+ });
+});
+
+describe("Ollama ChatLLM", () => {
+ const getInstance = () => {
+ return new OllamaChatLLM({
+ modelId: "llama3.1",
+ });
+ };
+
+ it("Serializes", async () => {
+ const instance = getInstance();
+ const serialized = instance.serialize();
+ const deserialized = OllamaChatLLM.fromSerialized(serialized);
+ verifyDeserialization(instance, deserialized);
+ });
+});
diff --git a/src/adapters/ollama/llm.ts b/src/adapters/ollama/llm.ts
new file mode 100644
index 00000000..6f5d24bb
--- /dev/null
+++ b/src/adapters/ollama/llm.ts
@@ -0,0 +1,190 @@
+/**
+ * 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 { LLM, LLMInput } from "@/llms/index.js";
+import { Emitter } from "@/emitter/emitter.js";
+import {
+ AsyncStream,
+ BaseLLMOutput,
+ BaseLLMTokenizeOutput,
+ ExecutionOptions,
+ GenerateCallbacks,
+ GenerateOptions,
+ LLMMeta,
+ LLMOutputError,
+ StreamGenerateOptions,
+} from "@/llms/base.js";
+import { GenerateResponse, Ollama as Client, Options as Parameters } from "ollama";
+import { GetRunContext } from "@/context.js";
+import { Cache } from "@/cache/decoratorCache.js";
+import { safeSum } from "@/internals/helpers/number.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { signalRace } from "@/internals/helpers/promise.js";
+import { customMerge } from "@/internals/helpers/object.js";
+import { extractModelMeta, registerClient } from "@/adapters/ollama/shared.js";
+
+interface Input {
+ modelId: string;
+ client?: Client;
+ parameters?: Partial;
+ executionOptions?: ExecutionOptions;
+}
+
+export class OllamaLLMOutput extends BaseLLMOutput {
+ public readonly results: GenerateResponse[];
+
+ constructor(result: GenerateResponse) {
+ super();
+ this.results = [result];
+ }
+
+ static {
+ this.register();
+ }
+
+ getTextContent(): string {
+ return this.finalResult.response;
+ }
+
+ @Cache()
+ get finalResult(): Readonly {
+ if (this.results.length === 0) {
+ throw new LLMOutputError("No chunks to get final result from!");
+ }
+
+ return customMerge(this.results, {
+ response: (value = "", oldValue = "") => oldValue + value,
+ total_duration: (value, oldValue) => value ?? oldValue,
+ load_duration: (value, oldValue) => value ?? oldValue,
+ model: (value, oldValue) => value ?? oldValue,
+ done: (value, oldValue) => value ?? oldValue,
+ done_reason: (value, oldValue) => value ?? oldValue,
+ created_at: (value, oldValue) => value ?? oldValue,
+ eval_duration: (value, oldValue) => value ?? oldValue,
+ prompt_eval_duration: (value, oldValue) => value ?? oldValue,
+ prompt_eval_count: safeSum,
+ eval_count: safeSum,
+ context: (value, oldValue) => [...(value || []), ...(oldValue || [])],
+ });
+ }
+
+ merge(other: OllamaLLMOutput): void {
+ Cache.getInstance(this, "finalResult").clear();
+ this.results.push(...other.results);
+ }
+
+ createSnapshot() {
+ return {
+ results: shallowCopy(this.results),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+}
+
+export class OllamaLLM extends LLM {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["ollama", "llm"],
+ creator: this,
+ });
+
+ public readonly client: Client;
+ public readonly parameters: Partial;
+
+ static {
+ this.register();
+ registerClient();
+ }
+
+ constructor({ client, modelId, parameters, executionOptions = {} }: Input) {
+ super(modelId, executionOptions);
+ this.client = client ?? new Client();
+ this.parameters = parameters ?? {};
+ }
+
+ protected async _generate(
+ input: LLMInput,
+ options: GenerateOptions,
+ run: GetRunContext,
+ ): Promise {
+ const response = await signalRace(
+ () =>
+ this.client.generate({
+ model: this.modelId,
+ stream: false,
+ raw: true,
+ prompt: input,
+ options: this.parameters,
+ format: options.guided?.json ? "json" : undefined,
+ }),
+ run.signal,
+ () => this.client.abort(),
+ );
+
+ return new OllamaLLMOutput(response);
+ }
+
+ protected async *_stream(
+ input: LLMInput,
+ options: StreamGenerateOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ for await (const chunk of await this.client.generate({
+ model: this.modelId,
+ stream: true,
+ raw: true,
+ prompt: input,
+ options: this.parameters,
+ format: options.guided?.json ? "json" : undefined,
+ })) {
+ if (run.signal.aborted) {
+ break;
+ }
+ yield new OllamaLLMOutput(chunk);
+ }
+ run.signal.throwIfAborted();
+ }
+
+ async meta(): Promise {
+ const model = await this.client.show({
+ model: this.modelId,
+ });
+
+ return extractModelMeta(model);
+ }
+
+ async tokenize(input: LLMInput): Promise {
+ return {
+ tokensCount: Math.ceil(input.length / 4),
+ };
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ modelId: this.modelId,
+ executionOptions: shallowCopy(this.executionOptions),
+ parameters: shallowCopy(this.parameters),
+ client: this.client,
+ };
+ }
+}
diff --git a/src/adapters/ollama/shared.ts b/src/adapters/ollama/shared.ts
new file mode 100644
index 00000000..c4d3113a
--- /dev/null
+++ b/src/adapters/ollama/shared.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 { Serializer } from "@/serializer/serializer.js";
+import { Config, Ollama as Client, ShowResponse } from "ollama";
+import { getPropStrict } from "@/internals/helpers/object.js";
+import { LLMMeta } from "@/llms/base.js";
+
+export function registerClient() {
+ Serializer.register(Client, {
+ toPlain: (value) => ({
+ config: getPropStrict(value, "config") as Config,
+ fetch: getPropStrict(value, "fetch"),
+ }),
+ fromPlain: (value) =>
+ new Client({
+ fetch: value.fetch ?? value.config.fetch,
+ host: value.config.host,
+ proxy: value.config.proxy,
+ }),
+ });
+}
+
+export function extractModelMeta(response: ShowResponse): LLMMeta {
+ const tokenLimit = Object.entries(response.model_info)
+ .find(([k]) => k.includes("context_length") || k.includes("max_sequence_length"))
+ ?.at(1);
+
+ return {
+ tokenLimit: tokenLimit || Infinity,
+ };
+}
diff --git a/src/adapters/openai/chat.test.ts b/src/adapters/openai/chat.test.ts
new file mode 100644
index 00000000..7d8e6387
--- /dev/null
+++ b/src/adapters/openai/chat.test.ts
@@ -0,0 +1,33 @@
+/**
+ * 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 { verifyDeserialization } from "@tests/e2e/utils.js";
+import { OpenAIChatLLM } from "@/adapters/openai/chat.js";
+
+describe("OpenAI ChatLLM", () => {
+ const getInstance = () => {
+ return new OpenAIChatLLM({
+ modelId: "gpt-4o",
+ });
+ };
+
+ it("Serializes", async () => {
+ const instance = getInstance();
+ const serialized = instance.serialize();
+ const deserialized = OpenAIChatLLM.fromSerialized(serialized);
+ verifyDeserialization(instance, deserialized);
+ });
+});
diff --git a/src/adapters/openai/chat.ts b/src/adapters/openai/chat.ts
new file mode 100644
index 00000000..3f4aa763
--- /dev/null
+++ b/src/adapters/openai/chat.ts
@@ -0,0 +1,241 @@
+/**
+ * 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 {
+ AsyncStream,
+ BaseLLMTokenizeOutput,
+ ExecutionOptions,
+ GenerateCallbacks,
+ GenerateOptions,
+ LLMMeta,
+ StreamGenerateOptions,
+} from "@/llms/base.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { ChatLLM, ChatLLMOutput } from "@/llms/chat.js";
+import { BaseMessage, RoleType } from "@/llms/primitives/message.js";
+import { Emitter } from "@/emitter/emitter.js";
+import { ClientOptions, OpenAI as Client } from "openai";
+import { GetRunContext } from "@/context.js";
+import { promptTokensEstimate } from "openai-chat-tokens";
+import { isString } from "remeda";
+import { Serializer } from "@/serializer/serializer.js";
+import { getPropStrict } from "@/internals/helpers/object.js";
+
+type Parameters = Omit;
+type Response = Omit;
+
+export class OpenAIChatLLMOutput extends ChatLLMOutput {
+ public readonly responses: Response[];
+
+ constructor(response: Response) {
+ super();
+ this.responses = [response];
+ }
+
+ static {
+ this.register();
+ }
+
+ get messages() {
+ return this.responses
+ .flatMap((response) => response.choices)
+ .flatMap((choice) =>
+ BaseMessage.of({
+ role: choice.delta.role as RoleType,
+ text: choice.delta.content!,
+ }),
+ );
+ }
+
+ getTextContent(): string {
+ return this.messages.map((msg) => msg.text).join("\n");
+ }
+
+ merge(other: OpenAIChatLLMOutput): void {
+ this.responses.push(...other.responses);
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+
+ createSnapshot() {
+ return {
+ responses: shallowCopy(this.responses),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType): void {
+ Object.assign(this, snapshot);
+ }
+}
+
+interface Input {
+ modelId?: Client.ChatModel;
+ client?: Client;
+ parameters?: Partial;
+ executionOptions?: ExecutionOptions;
+}
+
+export class OpenAIChatLLM extends ChatLLM {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["openai", "chat_llm"],
+ creator: this,
+ });
+
+ public readonly client: Client;
+ public readonly parameters: Partial;
+
+ constructor({ client, modelId = "gpt-4o", parameters, executionOptions = {} }: Input = {}) {
+ super(modelId, executionOptions);
+ this.client = client ?? new Client();
+ this.parameters = parameters ?? {};
+ }
+
+ static {
+ this.register();
+ Serializer.register(Client, {
+ toPlain: (value) => ({
+ options: getPropStrict(value, "_options") as ClientOptions,
+ }),
+ fromPlain: (value) => new Client(value.options),
+ });
+ }
+
+ async meta(): Promise {
+ if (
+ this.modelId.includes("gpt-4o") ||
+ this.modelId.includes("gpt-4-turbo") ||
+ this.modelId.includes("gpt-4-0125-preview") ||
+ this.modelId.includes("gpt-4-1106-preview")
+ ) {
+ return { tokenLimit: 128 * 1024 };
+ } else if (this.modelId.includes("gpt-4")) {
+ return { tokenLimit: 8 * 1024 };
+ } else if (this.modelId.includes("gpt-3.5-turbo")) {
+ return { tokenLimit: 16 * 1024 };
+ } else if (this.modelId.includes("gpt-3.5")) {
+ return { tokenLimit: 8 * 1024 };
+ }
+
+ return {
+ tokenLimit: Infinity,
+ };
+ }
+
+ async tokenize(input: BaseMessage[]): Promise {
+ const tokensCount = promptTokensEstimate({
+ messages: input.map(
+ (msg) =>
+ ({
+ role: msg.role,
+ content: msg.text,
+ }) as Client.Chat.ChatCompletionMessageParam,
+ ),
+ });
+
+ return {
+ tokensCount,
+ };
+ }
+
+ protected _prepareRequest(
+ input: BaseMessage[],
+ options: GenerateOptions,
+ ): Client.Chat.ChatCompletionCreateParams {
+ return {
+ ...this.parameters,
+ model: this.modelId,
+ stream: false,
+ messages: input.map(
+ (message) =>
+ ({
+ role: message.role,
+ content: message.text,
+ response_metadata: message.meta,
+ }) as Client.Chat.ChatCompletionMessageParam,
+ ),
+ ...(options?.guided?.json && {
+ response_format: {
+ type: "json_schema",
+ json_schema: isString(options.guided.json)
+ ? JSON.parse(options.guided.json)
+ : options.guided.json,
+ },
+ }),
+ };
+ }
+
+ protected async _generate(
+ input: BaseMessage[],
+ options: GenerateOptions,
+ run: GetRunContext,
+ ): Promise {
+ const response = await this.client.chat.completions.create(
+ {
+ ...this._prepareRequest(input, options),
+ stream: false,
+ },
+ {
+ signal: run.signal,
+ },
+ );
+
+ return new OpenAIChatLLMOutput({
+ id: response.id,
+ model: response.model,
+ created: response.created,
+ usage: response.usage,
+ service_tier: response.service_tier,
+ system_fingerprint: response.system_fingerprint,
+ choices: response.choices.map(
+ (choice) =>
+ ({
+ delta: choice.message,
+ index: 1,
+ logprobs: choice.logprobs,
+ finish_reason: choice.finish_reason,
+ }) as Client.Chat.ChatCompletionChunk.Choice,
+ ),
+ });
+ }
+
+ protected async *_stream(
+ input: BaseMessage[],
+ options: StreamGenerateOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ for await (const chunk of await this.client.chat.completions.create(
+ {
+ ...this._prepareRequest(input, options),
+ stream: true,
+ },
+ {
+ signal: run.signal,
+ },
+ )) {
+ yield new OpenAIChatLLMOutput(chunk);
+ }
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ parameters: shallowCopy(this.parameters),
+ client: this.client,
+ };
+ }
+}
diff --git a/src/adapters/watsonx/chat.ts b/src/adapters/watsonx/chat.ts
new file mode 100644
index 00000000..1fe24e72
--- /dev/null
+++ b/src/adapters/watsonx/chat.ts
@@ -0,0 +1,161 @@
+/**
+ * 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 { AsyncStream, GenerateCallbacks } from "@/llms/base.js";
+import {
+ WatsonXLLM,
+ WatsonXLLMGenerateOptions,
+ WatsonXLLMParameters,
+ WatsonXLLMOutput,
+} from "@/adapters/watsonx/llm.js";
+import { ChatLLM, ChatLLMOutput } from "@/llms/chat.js";
+import { BaseMessage, Role } from "@/llms/primitives/message.js";
+import { PromptTemplate } from "@/template.js";
+import { Cache } from "@/cache/decoratorCache.js";
+import { transformAsyncIterable } from "@/internals/helpers/stream.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { Emitter } from "@/emitter/emitter.js";
+import { GetRunContext } from "@/context.js";
+
+export class WatsonXChatLLMOutput extends ChatLLMOutput {
+ public readonly raw: WatsonXLLMOutput;
+
+ constructor(rawOutput: WatsonXLLMOutput) {
+ super();
+ this.raw = rawOutput;
+ }
+
+ @Cache()
+ get messages(): BaseMessage[] {
+ const text = this.raw.getTextContent();
+ return [
+ BaseMessage.of({
+ role: Role.ASSISTANT,
+ text,
+ meta: this.raw.meta,
+ }),
+ ];
+ }
+
+ merge(other: WatsonXChatLLMOutput): void {
+ Cache.getInstance(this, "messages").clear();
+ this.raw.merge(other.raw);
+ }
+
+ getTextContent(): string {
+ const [message] = this.messages;
+ return message.text;
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+
+ createSnapshot() {
+ return {
+ raw: shallowCopy(this.raw),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
+
+export interface WatsonXChatLLMInputConfig {
+ messagesToPrompt: PromptTemplate<"messages"> | ((messages: BaseMessage[]) => string);
+}
+
+export interface WatsonXChatLLMInput {
+ llm: WatsonXLLM;
+ config: WatsonXChatLLMInputConfig;
+}
+
+export class WatsonXChatLLM extends ChatLLM {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["watsonx", "chat_llm"],
+ creator: this,
+ });
+
+ public readonly llm: WatsonXLLM;
+ protected readonly config: WatsonXChatLLMInputConfig;
+ public readonly parameters: WatsonXLLMParameters;
+
+ constructor({ llm, config }: WatsonXChatLLMInput) {
+ super(llm.modelId, llm.executionOptions);
+ this.parameters = llm.parameters ?? {};
+ this.llm = llm;
+ this.config = config;
+ }
+
+ static {
+ this.register();
+ }
+
+ async meta() {
+ return this.llm.meta();
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ modelId: this.modelId,
+ parameters: this.parameters,
+ executionOptions: this.executionOptions,
+ llm: this.llm,
+ config: shallowCopy(this.config),
+ };
+ }
+
+ loadSnapshot(data: ReturnType): void {
+ super.loadSnapshot(data);
+ }
+
+ async tokenize(messages: BaseMessage[]) {
+ const prompt = this.messagesToPrompt(messages);
+ return this.llm.tokenize(prompt);
+ }
+
+ protected async _generate(
+ messages: BaseMessage[],
+ options: WatsonXLLMGenerateOptions,
+ run: GetRunContext,
+ ): Promise {
+ const prompt = this.messagesToPrompt(messages);
+ // @ts-expect-error protected property
+ const rawResponse = await this.llm._generate(prompt, options, run);
+ return new WatsonXChatLLMOutput(rawResponse);
+ }
+
+ protected async *_stream(
+ messages: BaseMessage[],
+ options: WatsonXLLMGenerateOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ const prompt = this.messagesToPrompt(messages);
+ // @ts-expect-error protected property
+ const response = this.llm._stream(prompt, options, run);
+ return yield* transformAsyncIterable(response, (output) => new WatsonXChatLLMOutput(output));
+ }
+
+ messagesToPrompt(messages: BaseMessage[]) {
+ const convertor = this.config.messagesToPrompt;
+ if (convertor instanceof PromptTemplate) {
+ return convertor.render({ messages });
+ }
+ return convertor(messages);
+ }
+}
diff --git a/src/adapters/watsonx/llm.ts b/src/adapters/watsonx/llm.ts
new file mode 100644
index 00000000..e6286b13
--- /dev/null
+++ b/src/adapters/watsonx/llm.ts
@@ -0,0 +1,413 @@
+/**
+ * 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 { LLM, LLMInput } from "@/llms/index.js";
+import {
+ AsyncStream,
+ BaseLLMOutput,
+ BaseLLMTokenizeOutput,
+ ExecutionOptions,
+ GenerateCallbacks,
+ GenerateOptions,
+ LLMError,
+ LLMFatalError,
+ LLMMeta,
+ LLMOutputError,
+} from "@/llms/base.js";
+import { HttpError } from "@ibm-generative-ai/node-sdk";
+import * as R from "remeda";
+import { FrameworkError } from "@/errors.js";
+import { Cache, CacheFn } from "@/cache/decoratorCache.js";
+import { shallowCopy } from "@/serializer/utils.js";
+import { safeSum } from "@/internals/helpers/number.js";
+import { omitUndefined } from "@/internals/helpers/object.js";
+import { createURLParams, RestfulClient, RestfulClientError } from "@/internals/fetcher.js";
+import { identity } from "remeda";
+import { Emitter } from "@/emitter/emitter.js";
+import { GetRunContext } from "@/context.js";
+
+export interface WatsonXLLMOutputMeta {
+ model_id: string;
+ created_at: string;
+}
+
+export interface WatsonXLLMOutputResult {
+ generated_text: string;
+ generated_token_count: number;
+ input_token_count: number;
+ stop_reason?: string;
+}
+
+export interface WatsonXLLMOutputConstructor {
+ meta: WatsonXLLMOutputMeta;
+ results: WatsonXLLMOutputResult[];
+ system: Record[]>;
+}
+
+export class WatsonXLLMOutput extends BaseLLMOutput {
+ public readonly meta: WatsonXLLMOutputMeta;
+ public readonly results: WatsonXLLMOutputResult[];
+
+ constructor(content: WatsonXLLMOutputConstructor) {
+ super();
+ this.meta = content.meta;
+ this.results = content.results;
+ }
+
+ static {
+ this.register();
+ }
+
+ getTextContent(): string {
+ return this.finalResult.generated_text;
+ }
+
+ @Cache()
+ get finalResult(): Readonly {
+ if (this.results.length === 0) {
+ throw new LLMOutputError("No chunks to get final result from!");
+ }
+
+ const processors: {
+ [K in keyof WatsonXLLMOutputResult]: (
+ value: WatsonXLLMOutputResult[K],
+ oldValue: WatsonXLLMOutputResult[K],
+ ) => WatsonXLLMOutputResult[K];
+ } = {
+ generated_text: (value = "", oldValue = "") => oldValue + value,
+ input_token_count: safeSum,
+ generated_token_count: safeSum,
+ stop_reason: (value, oldValue) => value ?? oldValue,
+ };
+
+ const finalResult = {} as WatsonXLLMOutputResult;
+ for (const next of this.results) {
+ for (const [key, value] of R.entries(next)) {
+ const oldValue = finalResult[key];
+ // @ts-expect-error weak typing due to generated types
+ finalResult[key] = (processors[key] ?? takeFirst)(value, oldValue);
+ }
+ }
+
+ return finalResult;
+ }
+
+ merge(other: WatsonXLLMOutput): void {
+ Cache.getInstance(this, "finalResult").clear();
+
+ this.results.push(...other.results);
+ Object.assign(this.meta, omitUndefined(other.meta));
+ }
+
+ createSnapshot() {
+ return {
+ results: shallowCopy(this.results),
+ meta: shallowCopy(this.meta),
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+
+ toString(): string {
+ return this.getTextContent();
+ }
+}
+
+export type WatsonXLLMParameters = Record;
+export type WatsonXLLMModerations = Record;
+
+export interface WatsonXLLMGenerateOptions extends GenerateOptions {
+ parameters?: WatsonXLLMParameters;
+ moderations?: WatsonXLLMModerations;
+}
+
+export interface WatsonXLLMInput {
+ modelId: string;
+ projectId?: string;
+ spaceId?: string;
+ deploymentId?: string;
+ version?: string;
+ apiKey?: string;
+ accessToken?: string;
+ baseUrl?: string;
+ authBaseUrl?: string;
+ region?: string;
+ parameters?: WatsonXLLMParameters;
+ moderations?: WatsonXLLMModerations;
+ executionOptions?: ExecutionOptions;
+ transform?: WatsonXLLMTransformFn;
+}
+
+type WatsonXLLMTransformFn = (body: Record) => Record;
+
+function createApiClient({
+ deploymentId,
+ apiKey,
+ baseUrl,
+ authBaseUrl = "https://iam.cloud.ibm.com",
+ region = "us-south",
+ accessToken,
+ version = "2023-05-02",
+ projectId,
+ spaceId,
+}: WatsonXLLMInput) {
+ const paths = (() => {
+ const pathPrefix = deploymentId ? `/ml/v1/deployments/${deploymentId}` : "/ml/v1";
+ const queryParams = createURLParams({
+ version,
+ });
+
+ return {
+ generate: `${pathPrefix}/text/generation?${queryParams}`,
+ generate_stream: `${pathPrefix}/text/generation_stream?${queryParams}`,
+ tokenization: `${pathPrefix}/text/tokenization?${queryParams}`,
+ models: `/ml/v1/foundation_model_specs?${queryParams}`,
+ deployment: deploymentId
+ ? `/ml/v4/deployments/${deploymentId}?${createURLParams({ version, project_id: projectId, space_id: projectId ? undefined : spaceId })}`
+ : "not_defined_endpoint",
+ };
+ })();
+
+ const getHeaders = CacheFn.create(async () => {
+ const getAccessToken = async () => {
+ if (accessToken) {
+ return { ttl: Infinity, token: accessToken };
+ }
+
+ const response = await fetch(new URL("/identity/token", authBaseUrl), {
+ method: "POST",
+ headers: {
+ "Content-Type": "application/x-www-form-urlencoded",
+ },
+ body: createURLParams({
+ grant_type: "urn:ibm:params:oauth:grant-type:apikey",
+ apikey: apiKey,
+ }),
+ });
+
+ if (!response.ok) {
+ throw new RestfulClientError("Failed to retrieve an API token.", [], {
+ context: response,
+ });
+ }
+
+ const data = await response.json();
+ if (!data?.access_token) {
+ throw new RestfulClientError("Access Token was not found in the response.");
+ }
+ return { ttl: (data.expires_in - 60) * 1000, token: data.access_token as string };
+ };
+
+ const response = await getAccessToken();
+ getHeaders.updateTTL(response.ttl);
+ return new Headers({
+ Authorization: `Bearer ${response.token}`,
+ Accept: "application/json",
+ "Content-Type": "application/json",
+ });
+ });
+
+ return new RestfulClient({
+ baseUrl: baseUrl || `https://${region}.ml.cloud.ibm.com`,
+ paths,
+ headers: getHeaders,
+ });
+}
+
+export class WatsonXLLM extends LLM {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["watsonx", "llm"],
+ creator: this,
+ });
+
+ protected client;
+ protected projectId;
+ protected deploymentId;
+ protected spaceId;
+ protected transform: WatsonXLLMTransformFn;
+ public readonly moderations;
+ public readonly parameters: WatsonXLLMParameters;
+
+ constructor(input: WatsonXLLMInput) {
+ super(input.modelId, input.executionOptions);
+ this.projectId = input.projectId;
+ this.spaceId = input.spaceId;
+ this.deploymentId = input.deploymentId;
+ this.moderations = input.moderations;
+ this.transform = input.transform ?? identity();
+ this.client = createApiClient(input);
+ this.parameters = input.parameters ?? {};
+ }
+
+ static {
+ this.register();
+ }
+
+ @Cache()
+ async meta(): Promise {
+ let modelId = this.modelId;
+ if (this.deploymentId) {
+ const { entity } = await this.client.fetch("deployment");
+ modelId = entity.base_model_id ?? modelId;
+ }
+
+ if (!modelId) {
+ throw new LLMFatalError(`Cannot retrieve metadata for model '${modelId ?? "undefined"}'`);
+ }
+
+ const {
+ resources: [model],
+ } = await this.client.fetch("models", {
+ searchParams: createURLParams({
+ filters: `modelid_${modelId}`,
+ limit: "1",
+ }),
+ });
+
+ return {
+ tokenLimit: model?.model_limits?.max_sequence_length ?? Infinity,
+ };
+ }
+
+ createSnapshot() {
+ return {
+ ...super.createSnapshot(),
+ modelId: this.modelId,
+ spaceId: this.spaceId,
+ deploymentId: this.deploymentId,
+ projectId: this.projectId,
+ parameters: shallowCopy(this.parameters),
+ moderations: shallowCopy(this.moderations),
+ executionOptions: shallowCopy(this.executionOptions),
+ transform: this.transform,
+ client: this.client,
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType): void {
+ super.loadSnapshot(snapshot);
+ }
+
+ protected _transformError(error: Error): Error {
+ if (error instanceof FrameworkError) {
+ throw error;
+ }
+ if (error instanceof HttpError) {
+ throw new LLMError("LLM has occurred an error!", [error], {
+ isRetryable: [408, 425, 429, 500, 503].includes(error.status_code),
+ });
+ }
+ return new LLMError("LLM has occurred an error!", [error]);
+ }
+
+ async tokenize(input: LLMInput): Promise {
+ try {
+ const { result } = await this.client.fetch("tokenization", {
+ method: "POST",
+ body: JSON.stringify({
+ input,
+ model_id: this.modelId,
+ parameters: {
+ return_tokens: true,
+ },
+ }),
+ });
+
+ return {
+ tokensCount: result.token_count,
+ tokens: result.tokens,
+ };
+ } catch (e) {
+ throw this._transformError(e);
+ }
+ }
+
+ protected async _generate(
+ input: LLMInput,
+ options: WatsonXLLMGenerateOptions,
+ run: GetRunContext,
+ ): Promise {
+ try {
+ const response = await this.client.fetch("generate", {
+ method: "POST",
+ body: JSON.stringify(
+ this.transform({
+ input,
+ ...(!this.deploymentId && {
+ model_id: this.modelId,
+ project_id: this.projectId,
+ space_id: this.projectId ? undefined : this.spaceId,
+ }),
+ parameters: options?.parameters ?? this.parameters,
+ moderations: options?.moderations ?? this.moderations,
+ }),
+ ),
+ signal: run.signal,
+ });
+ return this._rawResponseToOutput(response);
+ } catch (e) {
+ throw this._transformError(e);
+ }
+ }
+
+ protected async *_stream(
+ input: LLMInput,
+ options: WatsonXLLMGenerateOptions,
+ run: GetRunContext,
+ ): AsyncStream {
+ try {
+ const response = this.client.stream("generate_stream", {
+ method: "POST",
+ body: JSON.stringify(
+ this.transform({
+ input,
+ ...(!this.deploymentId && {
+ model_id: this.modelId,
+ project_id: this.projectId,
+ space_id: this.projectId ? undefined : this.spaceId,
+ }),
+ parameters: options?.parameters ?? this.parameters,
+ moderations: options?.moderations ?? this.moderations,
+ }),
+ ),
+ signal: run.signal,
+ });
+
+ for await (const msg of response) {
+ const content = JSON.parse(msg.data);
+ yield this._rawResponseToOutput(content);
+ }
+ } catch (e) {
+ throw this._transformError(e);
+ }
+ }
+
+ protected _rawResponseToOutput(raw: any) {
+ return new WatsonXLLMOutput({
+ results: raw.results ?? [],
+ meta: R.pickBy(
+ {
+ model_id: raw.model_id,
+ created_at: raw.created_at!,
+ },
+ R.isDefined,
+ ),
+ system: raw.system ?? [],
+ });
+ }
+}
diff --git a/src/agents/base.ts b/src/agents/base.ts
new file mode 100644
index 00000000..10fead3a
--- /dev/null
+++ b/src/agents/base.ts
@@ -0,0 +1,78 @@
+/**
+ * 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 { FrameworkError } from "@/errors.js";
+import { AgentMeta } from "@/agents/types.js";
+import { Serializable } from "@/internals/serializable.js";
+import { GetRunContext, GetRunInstance, Run, RunContext } from "@/context.js";
+import { Emitter } from "@/emitter/emitter.js";
+
+export class AgentError extends FrameworkError {}
+
+export interface BaseAgentRunOptions {
+ signal?: AbortSignal;
+}
+
+export abstract class BaseAgent<
+ TInput,
+ TOutput,
+ TOptions extends BaseAgentRunOptions = BaseAgentRunOptions,
+> extends Serializable {
+ private isRunning = false;
+
+ public abstract readonly emitter: Emitter;
+
+ public run(
+ ...[input, options]: Partial extends TOptions
+ ? [input: TInput, options?: TOptions]
+ : [input: TInput, options: TOptions]
+ ): Run> {
+ if (this.isRunning) {
+ throw new AgentError("Agent is already running!");
+ }
+
+ return RunContext.enter(
+ this,
+ async (context) => {
+ try {
+ // @ts-expect-error
+ return await this._run(input, options, context);
+ } catch (e) {
+ if (e instanceof AgentError) {
+ throw e;
+ } else {
+ throw new AgentError(`Error has occurred!`, [e]);
+ }
+ } finally {
+ this.isRunning = false;
+ }
+ },
+ options?.signal,
+ );
+ }
+
+ protected abstract _run(
+ input: TInput,
+ options: TOptions,
+ run: GetRunContext,
+ ): Promise;
+
+ destroy() {
+ this.emitter.destroy();
+ }
+
+ public abstract get meta(): AgentMeta;
+}
diff --git a/src/agents/bee/__snapshots__/parser.test.ts.snap b/src/agents/bee/__snapshots__/parser.test.ts.snap
new file mode 100644
index 00000000..37047d29
--- /dev/null
+++ b/src/agents/bee/__snapshots__/parser.test.ts.snap
@@ -0,0 +1,26 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`Bee Parser > Chunking > Text 1`] = `
+{
+ "final_answer": "I need to search Colorado, find
+
+ the
+ area that the eastern sector of the Colorado extends into, then find the elevation range of the area.
+
+
+
+Extra Content.",
+}
+`;
+
+exports[`Bee Parser > Chunking > Text 2`] = `
+{
+ "final_answer": "ABC",
+}
+`;
+
+exports[`ReAct Parser > Chunking > Text 2`] = `
+{
+ "final_answer": "ABC",
+}
+`;
diff --git a/src/agents/bee/agent.ts b/src/agents/bee/agent.ts
new file mode 100644
index 00000000..ad7369a8
--- /dev/null
+++ b/src/agents/bee/agent.ts
@@ -0,0 +1,176 @@
+/**
+ * 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 { BaseAgent } from "@/agents/base.js";
+import { AnyTool } from "@/tools/base.js";
+import { BaseMemory } from "@/memory/base.js";
+import { ChatLLM, ChatLLMOutput } from "@/llms/chat.js";
+import { BaseMessage, Role } from "@/llms/primitives/message.js";
+import { AgentMeta } from "@/agents/types.js";
+import { BeeAgentTemplate, BeeAssistantPrompt } from "@/agents/bee/prompts.js";
+import * as R from "remeda";
+import { Emitter } from "@/emitter/emitter.js";
+import {
+ BeeAgentRunIteration,
+ BeeCallbacks,
+ BeeRunInput,
+ BeeRunOptions,
+ BeeRunOutput,
+} from "@/agents/bee/types.js";
+import { GetRunContext } from "@/context.js";
+import { BeeAgentRunner } from "@/agents/bee/runner.js";
+import { BeeAgentError } from "@/agents/bee/errors.js";
+import { BeeIterationToolResult } from "@/agents/bee/parser.js";
+
+export interface BeeInput {
+ llm: ChatLLM;
+ tools: AnyTool[];
+ memory: BaseMemory;
+ promptTemplate?: BeeAgentTemplate;
+ meta?: AgentMeta;
+}
+
+export class BeeAgent extends BaseAgent {
+ public readonly emitter = Emitter.root.child({
+ namespace: ["agent", "bee"],
+ creator: this,
+ });
+
+ constructor(protected readonly input: BeeInput) {
+ super();
+
+ const duplicate = input.tools.find((a, i, arr) =>
+ arr.find((b, j) => i !== j && a.name.toUpperCase() === b.name.toUpperCase()),
+ );
+ if (duplicate) {
+ throw new BeeAgentError(
+ `Agent's tools must all have different names. Conflicting tool: ${duplicate.name}.`,
+ );
+ }
+ }
+
+ static {
+ this.register();
+ }
+
+ get meta(): AgentMeta {
+ if (this.input.meta) {
+ return this.input.meta;
+ }
+
+ const tools = this.input.tools;
+
+ return {
+ name: "Bee",
+ description:
+ "The Bee framework demonstrates its ability to auto-correct and adapt in real-time, improving the overall reliability and resilience of the system.",
+ ...(tools.length > 0 && {
+ extraDescription: [
+ `Tools that I can use to accomplish given task.`,
+ ...tools.map((tool) => `Tool '${tool.name}': ${tool.description}.`),
+ ].join("\n"),
+ }),
+ };
+ }
+
+ protected async _run(
+ input: BeeRunInput,
+ options: BeeRunOptions = {},
+ run: GetRunContext,
+ ): Promise {
+ const iterations: BeeAgentRunIteration[] = [];
+ const maxIterations = options?.execution?.maxIterations ?? Infinity;
+
+ const runner = await BeeAgentRunner.create(run, this.input, options, input.prompt);
+
+ let finalMessage: BaseMessage | undefined;
+ while (!finalMessage) {
+ if (iterations.length >= maxIterations) {
+ throw new BeeAgentError(
+ `Agent was not able to resolve the task in ${maxIterations} iterations.`,
+ [],
+ { isFatal: true },
+ );
+ }
+
+ const iteration = await runner.llm();
+
+ if (iteration.state.tool_name || iteration.state.tool_caption || iteration.state.tool_input) {
+ const { output, success } = await runner.tool(iteration.state as BeeIterationToolResult);
+
+ for (const key of ["partialUpdate", "update"] as const) {
+ await run.emitter.emit(key, {
+ data: {
+ tool_output: output,
+ },
+ update: { key: "tool_output", value: output },
+ meta: { success },
+ });
+ }
+
+ await runner.memory.add(
+ BaseMessage.of({
+ role: Role.ASSISTANT,
+ text: BeeAssistantPrompt.clone().render({
+ toolName: [iteration.state.tool_name].filter(R.isTruthy),
+ toolCaption: [iteration.state.tool_caption].filter(R.isTruthy),
+ toolInput: [iteration.state.tool_input]
+ .filter(R.isTruthy)
+ .map((call) => JSON.stringify(call)),
+ thought: [iteration.state.thought].filter(R.isTruthy),
+ finalAnswer: [iteration.state.final_answer].filter(R.isTruthy),
+ toolOutput: [output],
+ }),
+ meta: { success },
+ }),
+ );
+
+ iteration.state.tool_output = output;
+ }
+ if (iteration.state.final_answer) {
+ finalMessage = BaseMessage.of({
+ role: Role.ASSISTANT,
+ text: iteration.state.final_answer,
+ });
+ await run.emitter.emit("success", {
+ data: finalMessage,
+ });
+ await runner.memory.add(finalMessage);
+ }
+ iterations.push(iteration);
+ }
+
+ await this.input.memory.addMany([
+ BaseMessage.of({
+ role: Role.USER,
+ text: input.prompt,
+ }),
+ finalMessage,
+ ]);
+ return { result: finalMessage, iterations, memory: runner.memory };
+ }
+
+ createSnapshot() {
+ return {
+ input: this.input,
+ emitter: this.emitter,
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
diff --git a/src/agents/bee/errors.ts b/src/agents/bee/errors.ts
new file mode 100644
index 00000000..9dcadd34
--- /dev/null
+++ b/src/agents/bee/errors.ts
@@ -0,0 +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 { AgentError } from "@/agents/base.js";
+
+export class BeeAgentError extends AgentError {}
diff --git a/src/agents/bee/parser.test.ts b/src/agents/bee/parser.test.ts
new file mode 100644
index 00000000..c5bc5118
--- /dev/null
+++ b/src/agents/bee/parser.test.ts
@@ -0,0 +1,182 @@
+/**
+ * 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 { BeeOutputParser, BeeOutputParserError } from "@/agents/bee/parser.js";
+import * as R from "remeda";
+
+import { omitEmptyValues } from "@/internals/helpers/object.js";
+import { expect } from "vitest";
+
+describe("Bee Parser", () => {
+ describe("Parsing", () => {
+ it("Basics", async () => {
+ const parser = new BeeOutputParser({});
+ await parser.add("Final Answer: I need to find the current president of the Czech Republic.");
+ await parser.finalize();
+ parser.validate();
+
+ const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy));
+ expect(result).toMatchInlineSnapshot(`
+ {
+ "final_answer": "I need to find the current president of the Czech Republic.",
+ }
+ `);
+ });
+
+ it("Ends up with the same result", async () => {
+ const partial = new Map();
+ const final = new Map();
+
+ const parser = new BeeOutputParser();
+ parser.emitter.on("update", async ({ update, type }) => {
+ if (type === "full") {
+ final.set(update.key, update.value);
+ } else {
+ partial.set(update.key, (partial.get(update.key) ?? "").concat(update.value));
+ }
+ });
+ await parser.add("Thought: I ");
+ await parser.add("will do it.");
+ await parser.finalize();
+ parser.validate();
+ expect(partial).toStrictEqual(final);
+ });
+
+ it("Parses chunked JSON", async () => {
+ const parser = new BeeOutputParser({});
+ await parser.add("Tool Name:\n");
+ await parser.add("Goo");
+ await parser.add("gle");
+ await parser.add("\n");
+ await parser.add("Tool ");
+ await parser.add("Input:\n");
+ await parser.add('{"query');
+ await parser.add('": "Czech President"');
+ await parser.add("}");
+ await parser.finalize();
+ parser.validate();
+
+ const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy));
+ expect(result).toMatchInlineSnapshot(`
+ {
+ "tool_input": {
+ "query": "Czech President",
+ },
+ "tool_name": "Google",
+ }
+ `);
+ });
+
+ it("Handles newlines", async () => {
+ const parser = new BeeOutputParser({
+ allowMultiLines: true,
+ preserveNewLines: false,
+ trimContent: true,
+ });
+ await parser.add("Final Answer:\n\nI need to find\n\n the fastest car. \n ");
+ await parser.finalize();
+ parser.validate();
+
+ const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy));
+ expect(result).toMatchInlineSnapshot(`
+ {
+ "final_answer": "I need to find the fastest car.",
+ }
+ `);
+ });
+
+ it("Ignores newlines before first keyword occurrence", async () => {
+ const parser = new BeeOutputParser();
+ await parser.add("");
+ await parser.add(" \n");
+ await parser.add(" \n ");
+ await parser.add("");
+ await parser.add("\n Final Answer: Hello");
+ await parser.finalize();
+ parser.validate();
+
+ const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy));
+ expect(result).toMatchInlineSnapshot(`
+ {
+ "final_answer": "Hello",
+ }
+ `);
+ });
+
+ it("Handles newlines with custom settings", async () => {
+ const parser = new BeeOutputParser({
+ allowMultiLines: true,
+ preserveNewLines: true,
+ trimContent: true,
+ });
+ await parser.add("Final Answer:\n\nI need to find\n\n the fastest car. \n ");
+ await parser.finalize();
+ parser.validate();
+
+ const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy));
+ expect(result).toMatchInlineSnapshot(`
+ {
+ "final_answer": "I need to find
+ the fastest car.",
+ }
+ `);
+ });
+ });
+
+ describe("Chunking", () => {
+ it.each([
+ " F#inal #answer : #I need to# search #Colorado, find#\n#\n the\n #area th#at th#e easter#n secto#r of# the Colora#do ex#tends i#nto, then find th#e elev#ation# #range #of the area.\n\n\n\nExtra Content.",
+ "\nfinal answer:A#B#C###",
+ ])("Text", async (text) => {
+ const parser = new BeeOutputParser({
+ allowMultiLines: true,
+ preserveNewLines: true,
+ trimContent: false,
+ });
+
+ const chunks = text.split("#");
+ for (const chunk of chunks) {
+ await parser.add(chunk);
+ }
+
+ await parser.finalize();
+ parser.validate();
+
+ const result = R.pipe(parser.parse(), R.pickBy(R.isTruthy));
+ expect(result).toMatchSnapshot();
+ });
+ });
+
+ describe("Validation", () => {
+ it("Throws when no data passed", () => {
+ const parser = new BeeOutputParser();
+ expect(omitEmptyValues(parser.parse())).toStrictEqual({});
+ expect(() => parser.validate()).toThrowError(BeeOutputParserError);
+ });
+
+ it.each(["Hello\nWorld", "Tool{\nxxx", "\n\n\nT"])(
+ "Throws on invalid data (%s)",
+ async (chunk) => {
+ const parser = new BeeOutputParser({
+ allowMultiLines: false,
+ });
+ await expect(parser.add(chunk)).rejects.toThrowError(BeeOutputParserError);
+ await expect(parser.finalize()).rejects.toThrowError(BeeOutputParserError);
+ expect(() => parser.validate()).toThrowError(BeeOutputParserError);
+ },
+ );
+ });
+});
diff --git a/src/agents/bee/parser.ts b/src/agents/bee/parser.ts
new file mode 100644
index 00000000..10acd235
--- /dev/null
+++ b/src/agents/bee/parser.ts
@@ -0,0 +1,293 @@
+/**
+ * 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 * as R from "remeda";
+import { FrameworkError } from "@/errors.js";
+import { Cache } from "@/cache/decoratorCache.js";
+import { halveString } from "@/internals/helpers/string.js";
+import { parseBrokenJson } from "@/internals/helpers/schema.js";
+import { Emitter } from "@/emitter/emitter.js";
+import { NonUndefined } from "@/internals/types.js";
+
+export interface BeeOutputParserOptions {
+ allowMultiLines: boolean;
+ preserveNewLines: boolean;
+ trimContent: boolean;
+}
+
+interface BeeIterationSerializedResult {
+ thought?: string;
+ tool_name?: string;
+ tool_caption?: string;
+ tool_input?: string;
+ tool_output?: string;
+ final_answer?: string;
+}
+
+export interface BeeIterationResult extends Omit {
+ tool_input?: unknown;
+}
+
+export type BeeIterationToolResult = NonUndefined;
+
+export class BeeOutputParserError extends FrameworkError {}
+
+const NEW_LINE_CHARACTER = "\n";
+
+interface Callbacks {
+ update: (data: {
+ type: "full" | "partial";
+ state: BeeIterationResult;
+ update: {
+ key: keyof BeeIterationResult;
+ value: string;
+ };
+ }) => Promise;
+}
+
+export class BeeOutputParser {
+ public isDone = false;
+ protected readonly lines: string[];
+ protected lastKeyModified: keyof BeeIterationSerializedResult | null = null;
+ public stash: string;
+ public readonly emitter = new Emitter({
+ creator: this,
+ namespace: ["agent", "bee", "parser"],
+ });
+
+ protected readonly options: BeeOutputParserOptions;
+
+ protected readonly result: BeeIterationSerializedResult = {
+ thought: undefined,
+ tool_name: undefined,
+ tool_caption: undefined,
+ tool_input: undefined,
+ tool_output: undefined,
+ final_answer: undefined,
+ };
+
+ constructor(options?: Partial) {
+ this.options = {
+ ...this._defaultOptions,
+ ...options,
+ };
+ this.lines = [];
+ this.stash = "";
+ }
+
+ async add(chunk: string) {
+ chunk = chunk ?? "";
+ this.stash += chunk;
+ if (!chunk.includes(NEW_LINE_CHARACTER)) {
+ return;
+ }
+
+ while (this.stash.includes(NEW_LINE_CHARACTER)) {
+ this.filterStash();
+ const [line, stash = ""] = halveString(this.stash, NEW_LINE_CHARACTER);
+ this.stash = stash;
+
+ await this._processChunk(line);
+ }
+ }
+
+ protected filterStash() {
+ this.stash = this.stash.replaceAll("<|eom_id|>", "");
+ this.stash = this.stash.replaceAll("<|eot_id|>", "");
+ this.stash = this.stash.replaceAll("<|start_header_id|>assistant<|end_header_id|>", "");
+ this.stash = this.stash.replaceAll("<|start_header_id|>", "");
+ this.stash = this.stash.replaceAll("<|end_header_id|>", "");
+ this.stash = this.stash.replaceAll("<|im_start|>", "");
+ this.stash = this.stash.replaceAll("<|im_end|>", "");
+ }
+
+ async finalize() {
+ if (this.stash) {
+ await this._processChunk(this.stash);
+ this.stash = "";
+ }
+
+ if (this.isEmpty()) {
+ const response = this.lines.join(NEW_LINE_CHARACTER).concat(this.stash);
+ this.lines.length = 0;
+ this.stash = "";
+
+ await this.add(`Thought: ${response}${NEW_LINE_CHARACTER}`);
+ await this.add(`Final Answer: ${response}${NEW_LINE_CHARACTER}`);
+ }
+ if (this.result.thought && !this.result.final_answer && !this.result.tool_input) {
+ this.stash = "";
+ await this.add(`Final Answer: ${this.result.thought}${NEW_LINE_CHARACTER}`);
+ }
+
+ if (this.lastKeyModified) {
+ const parsed = this.parse();
+ await this.emitter.emit("update", {
+ type: "full",
+ state: parsed,
+ update: {
+ key: this.lastKeyModified,
+ value: this.result[this.lastKeyModified]!,
+ },
+ });
+ }
+ this.lastKeyModified = null;
+ }
+
+ isEmpty() {
+ return R.isEmpty(R.pickBy(this.result, R.isTruthy));
+ }
+
+ validate() {
+ if (this.isEmpty()) {
+ throw new BeeOutputParserError("Nothing valid has been parsed yet!", [], {
+ context: {
+ raw: this.lines.join(NEW_LINE_CHARACTER),
+ stash: this.stash,
+ },
+ });
+ }
+
+ const { final_answer, tool_name, tool_input } = this.parse();
+ const context = {
+ result: this.parse(),
+ stash: this.stash,
+ };
+ if (!final_answer && !tool_input) {
+ if (this.result.tool_input) {
+ throw new BeeOutputParserError('Invalid "Tool Input" has been generated.', [], {
+ context: {
+ toolName: tool_name,
+ toolCaption: this.result.tool_caption,
+ toolInput: this.result.tool_input,
+ ...context,
+ },
+ });
+ }
+
+ throw new BeeOutputParserError('Neither "Final Answer" nor "Tool Call" are present.', [], {
+ context,
+ });
+ }
+ if (tool_input && final_answer) {
+ throw new BeeOutputParserError('Both "Final Answer" and "Tool Call" are present.', [], {
+ context,
+ });
+ }
+ }
+
+ @Cache()
+ protected get _defaultOptions(): BeeOutputParserOptions {
+ return {
+ allowMultiLines: true,
+ preserveNewLines: true,
+ trimContent: false,
+ };
+ }
+
+ protected async _processChunk(chunk: string) {
+ if (this.isDone) {
+ return;
+ }
+
+ this.lines.push(chunk);
+
+ let oldChunk = this.lastKeyModified ? this.result[this.lastKeyModified] : "";
+ if (!this._extractStepPair(chunk) && this.options.allowMultiLines && this.lastKeyModified) {
+ const prev = this.result[this.lastKeyModified] || "";
+ const newLine = this.options.preserveNewLines ? NEW_LINE_CHARACTER : "";
+ chunk = `${this.lastKeyModified}:${prev}${newLine}${chunk}`;
+ }
+
+ const step = this._extractStepPair(chunk);
+ if (!step && this.lastKeyModified === null && this.options.allowMultiLines) {
+ return;
+ }
+
+ if (!step) {
+ throw new BeeOutputParserError(`No valid type has been detected in the chunk. (${chunk})}`);
+ }
+
+ if (this.lastKeyModified && this.lastKeyModified !== step.type) {
+ this.isDone = Boolean(this.result[step.type]);
+ if (this.isDone) {
+ return;
+ }
+
+ const state = this.parse();
+ await this.emitter.emit("update", {
+ type: "full",
+ state,
+ update: {
+ key: this.lastKeyModified,
+ value: this.result[this.lastKeyModified]!,
+ },
+ });
+ oldChunk = this.result[step.type] ?? "";
+ }
+ this.lastKeyModified = step.type;
+
+ if (step.content) {
+ this.result[step.type] = step.content;
+ const state = this.parse();
+ await this.emitter.emit("update", {
+ type: "partial",
+ state,
+ update: {
+ key: step.type,
+ value: step.content.replace(oldChunk ?? "", ""),
+ },
+ });
+ }
+ }
+
+ parse(): BeeIterationResult {
+ const toolInput = parseBrokenJson(this?.result.tool_input?.trim?.(), { pair: ["{", "}"] });
+ return R.pickBy(
+ Object.assign(
+ { ...this.result },
+ {
+ tool_name: this.result.tool_name,
+ tool_input: toolInput ?? undefined,
+ },
+ ),
+ R.isDefined,
+ );
+ }
+
+ protected _isValidStepType(type?: string | null): type is keyof BeeIterationSerializedResult {
+ return Boolean(type && type in this.result);
+ }
+
+ protected _extractStepPair(line: string) {
+ let [, type, content] = line.match(/\s*([\w|\s]+?)\s*:\s*(.*)/ms) ?? [line, null, null];
+ type = type ? type.trim().toLowerCase().replace(" ", "_") : null;
+
+ if (!this._isValidStepType(type)) {
+ return null;
+ }
+
+ content = content ?? "";
+ if (this.options.trimContent) {
+ content = content.trim();
+ }
+
+ return {
+ type,
+ content,
+ };
+ }
+}
diff --git a/src/agents/bee/prompts.ts b/src/agents/bee/prompts.ts
new file mode 100644
index 00000000..89d44b46
--- /dev/null
+++ b/src/agents/bee/prompts.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 { PromptTemplate } from "@/template.js";
+
+export const BeeAgentSystemPrompt = new PromptTemplate({
+ variables: ["instructions", "tools", "tool_names"] as const,
+ defaults: {
+ instructions: "You are a helpful assistant that uses tools to answer questions.",
+ },
+ template: `{{instructions}}
+
+# Tools
+
+Tools must be used to retrieve factual or historical information to answer the question.
+{{#tools.length}}
+A tool can be used by generating the following three lines:
+
+Tool Name: ZblorgColorLookup
+Tool Caption: Searching Zblorg #178
+Tool Input: {"id":178}
+
+## Available tools
+{{#tools}}
+Tool Name: {{name}}
+Tool Description: {{description}}
+Tool Input Schema: {{schema}}
+
+{{/tools}}
+{{/tools.length}}
+{{^tools.length}}
+## Available tools
+
+No tools are available at the moment therefore you mustn't provide any factual or historical information.
+If you need to, you must respond that you cannot.
+{{/tools.length}}
+
+# Instructions
+
+Responses must always have the following structure:
+- The user's input starts with 'Question: ' followed by the question the user asked, for example, 'Question: What is the color of Zblorg #178?'
+ - The question may contain square brackets with a nested sentence, like 'What is the color of [The Zblorg with the highest score of the 2023 season is Zblorg #178.]?'. Just assume that the question regards the entity described in the bracketed sentence, in this case 'Zblorg #178'.
+- Line starting 'Thought: ', explaining the thought, for example 'Thought: I don't know what Zblorg is, but given that I have a ZblorgColorLookup tool, I can assume that it is something that can have a color and I should use the ZblorgColorLookup tool to find out the color of Zblorg number 178.'
+ - In a 'Thought', it is either determined that a Tool Call will be performed to obtain more information, or that the available information is sufficient to provide the Final Answer.
+ - If a tool needs to be called and is available, the following lines will be:
+ - Line starting 'Tool Name: ' name of the tool that you want to use.
+ - Line starting 'Tool Caption: ' short description of the calling action.
+ - Line starting 'Tool Input: ' JSON formatted input adhering to the selected tool JSON Schema.
+ - Line starting 'Tool Output: ', containing the tool output, for example 'Tool Output: {"success": true, "color": "green"}'
+ - The 'Tool Output' may or may not bring useful information. The following 'Thought' must determine whether the information is relevant and how to proceed further.
+ - If enough information is available to provide the Final Answer, the following line will be:
+ - Line starting 'Final Answer: ' followed by a response to the original question and context, for example: 'Final Answer: Zblorg #178 is green.'
+ - Use markdown syntax for formatting code snippets, links, JSON, tables, images, files.
+ - To reference an internal file, use the markdown syntax [file_name.ext](urn:file_identifier).
+ - The bracketed part must contain the file name, verbatim.
+ - The parenthesis part must contain the file URN, which can be obtained from the user or from tools.
+ - The agent does not, under any circumstances, reference a URN that was not provided by the user or a tool in the current conversation.
+ - To show an image, prepend an exclamation mark, as usual in markdown: ![file_name.ext](urn:file_identifier).
+ - This only applies to internal files. HTTP(S) links must be provided as is, without any modifications.
+- The sequence of lines will be 'Thought' - ['Tool Name' - 'Tool Caption' - 'Tool Input' - 'Tool Output' - 'Thought'] - 'Final Answer', with the bracketed part repeating one or more times (but never repeating them in a row). Do not use empty lines between instructions.
+- Sometimes, things don't go as planned. Tools may not provide useful information on the first few tries. The agent always tries a few different approaches before declaring the problem unsolvable:
+- When the tool doesn't give you what you were asking for, you MUST either use another tool or a different tool input.
+ - When using search engines, the assistant tries different formulations of the query, possibly even in a different language.
+- When executing code, the assistant fixes and retries when the execution errors out and tries a completely different approach if the code does not seem to be working.
+ - When the problem seems too hard for the tool, the assistant tries to split the problem into a few smaller ones.
+
+## Notes
+- Any comparison table (including its content), file, image, link, or other asset must only be in the Final Answer.
+- When the question is unclear, respond with a line starting with 'Final Answer:' followed by the information needed to solve the problem.
+- When the user wants to chitchat instead, always respond politely.
+- IMPORTANT: Lines 'Thought', 'Tool Name', 'Tool Caption', 'Tool Input', 'Tool Output' and 'Final Answer' must be sent within a single message.
+`,
+});
+export type BeeAgentTemplate = typeof BeeAgentSystemPrompt;
+
+export const BeeAssistantPrompt = new PromptTemplate({
+ variables: ["thought", "toolName", "toolCaption", "toolInput", "toolOutput", "finalAnswer"],
+ optionals: ["thought", "toolName", "toolCaption", "toolInput", "toolOutput", "finalAnswer"],
+ template: `{{#thought}}Thought: {{.}}\n{{/thought}}{{#toolName}}Tool Name: {{.}}\n{{/toolName}}{{#toolCaption}}Tool Caption: {{.}}\n{{/toolCaption}}{{#toolInput}}Tool Input: {{.}}\n{{/toolInput}}{{#toolOutput}}Tool Output: {{.}}\n{{/toolOutput}}{{#finalAnswer}}Final Answer: {{.}}{{/finalAnswer}}`,
+});
+
+export const BeeUserPrompt = new PromptTemplate({
+ variables: ["input"],
+ template: `Question: {{input}}`,
+});
+
+export const BeeToolErrorPrompt = new PromptTemplate({
+ variables: ["reason"],
+ template: `The tool has failed; the error log is shown below. If the tool cannot accomplish what you want, use a different tool or explain why you can't use it.
+
+{{reason}}`,
+});
+
+export const BeeToolInputErrorPrompt = new PromptTemplate({
+ variables: ["reason"],
+ template: `{{reason}}
+
+HINT: If you're convinced that the input was correct but the tool cannot process it then use a different tool or say I don't know.`,
+});
+
+export const BeeToolNoResultsPrompt = new PromptTemplate({
+ variables: [],
+ template: `No results were found!`,
+});
diff --git a/src/agents/bee/runner.ts b/src/agents/bee/runner.ts
new file mode 100644
index 00000000..8a102ea2
--- /dev/null
+++ b/src/agents/bee/runner.ts
@@ -0,0 +1,297 @@
+/**
+ * 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 { BeeAgentError } from "@/agents/bee/errors.js";
+
+import { BaseMessage, Role } from "@/llms/primitives/message.js";
+import { TokenMemory } from "@/memory/tokenMemory.js";
+import { PromptTemplate } from "@/template.js";
+import { BeeAgentRunIteration, BeeCallbacks, BeeRunOptions } from "@/agents/bee/types.js";
+import { RequiredAll } from "@/internals/types.js";
+import { BaseToolRunOptions, ToolInputValidationError, ToolOutput } from "@/tools/base.js";
+import { getProp } from "@/internals/helpers/object.js";
+import { Retryable } from "@/internals/helpers/retryable.js";
+import { FrameworkError } from "@/errors.js";
+import { RunContext } from "@/context.js";
+import { BeeInput } from "@/agents/bee/agent.js";
+import { RetryCounter } from "@/internals/helpers/counter.js";
+import {
+ BeeAgentSystemPrompt,
+ BeeToolErrorPrompt,
+ BeeToolInputErrorPrompt,
+ BeeToolNoResultsPrompt,
+ BeeUserPrompt,
+} from "@/agents/bee/prompts.js";
+import { BeeIterationToolResult, BeeOutputParser } from "@/agents/bee/parser.js";
+
+export class BeeAgentRunnerFatalError extends BeeAgentError {
+ isFatal = true;
+}
+
+export class BeeAgentRunner {
+ protected readonly failedAttemptsCounter;
+
+ constructor(
+ protected readonly run: RunContext>,
+ protected readonly input: BeeInput,
+ protected readonly options: BeeRunOptions,
+ public readonly memory: TokenMemory,
+ ) {
+ this.failedAttemptsCounter = new RetryCounter(options?.execution?.totalMaxRetries);
+ }
+
+ static async create(
+ run: RunContext>,
+ input: BeeInput,
+ options: BeeRunOptions,
+ prompt: string,
+ ) {
+ const memory = new TokenMemory({
+ llm: input.llm,
+ capacityThreshold: 0.85,
+ handlers: {
+ removalSelector(curMessages) {
+ // First we remove messages from the past conversations
+ const prevConversationMessage = curMessages.find((msg) =>
+ input.memory.messages.includes(msg),
+ );
+ if (prevConversationMessage) {
+ return prevConversationMessage;
+ }
+
+ if (curMessages.length <= 3) {
+ throw new BeeAgentRunnerFatalError(
+ "Cannot fit the current conversation into the context window!",
+ );
+ }
+
+ const lastMessage =
+ curMessages.find(
+ (msg) => msg.role === Role.ASSISTANT && getProp(msg, ["ctx", "success"]) === false,
+ ) ?? curMessages.find((msg) => msg.role === Role.ASSISTANT);
+
+ if (!lastMessage) {
+ throw new BeeAgentRunnerFatalError(
+ "Cannot fit the current conversation into the context window!",
+ );
+ }
+ return lastMessage;
+ },
+ },
+ });
+ const template = input.promptTemplate ?? BeeAgentSystemPrompt;
+ await memory.addMany([
+ BaseMessage.of({
+ role: Role.SYSTEM,
+ text: template.render({
+ tools: await Promise.all(
+ input.tools.map(async (tool) => ({
+ name: tool.name,
+ description: tool.description.replaceAll("\n", ".").replace(/\.$/, "").concat("."),
+ schema: JSON.stringify(await tool.getInputJsonSchema()),
+ })),
+ ),
+ tool_names: input.tools.map((tool) => tool.name).join(","),
+ instructions: PromptTemplate.defaultPlaceholder,
+ }),
+ }),
+ ...input.memory.messages,
+ BaseMessage.of({
+ role: Role.USER,
+ text: BeeUserPrompt.clone().render({
+ input: prompt.trim() ? prompt : "Empty message.",
+ }),
+ }),
+ ]);
+
+ return new BeeAgentRunner(run, input, options, memory);
+ }
+
+ async llm(): Promise {
+ const emitter = this.run.emitter;
+
+ return new Retryable({
+ onRetry: () => emitter.emit("retry", undefined),
+ onError: async (error) => {
+ await emitter.emit("error", { error });
+ this.failedAttemptsCounter.use(error);
+ },
+ executor: async () => {
+ await emitter.emit("start", undefined);
+
+ const outputParser = new BeeOutputParser();
+ const llmOutput = await this.input.llm
+ .generate(this.memory.messages.slice(), {
+ signal: this.run.signal,
+ stream: true,
+ guided: {
+ regex:
+ /Thought:.+\n(?:Final Answer:[\S\s]+|Tool Name:.+\nTool Caption:.+\nTool Input:\{.+\}\nTool Output:)/
+ .source,
+ },
+ })
+ .observe((llmEmitter) => {
+ outputParser.emitter.on("update", async ({ type, update, state }) => {
+ await emitter.emit(type === "full" ? "update" : "partialUpdate", {
+ data: state,
+ update,
+ meta: { success: true },
+ });
+ });
+
+ llmEmitter.on("newToken", async ({ value, callbacks }) => {
+ if (outputParser.isDone) {
+ callbacks.abort();
+ return;
+ }
+
+ await outputParser.add(value.getTextContent());
+ if (outputParser.stash.match(/^\s*Tool Output:/i)) {
+ outputParser.stash = "";
+ callbacks.abort();
+ }
+ });
+ });
+
+ await outputParser.finalize();
+ outputParser.validate();
+
+ return {
+ state: outputParser.parse(),
+ raw: llmOutput,
+ };
+ },
+ config: {
+ maxRetries: this.options.execution?.maxRetriesPerStep,
+ signal: this.run.signal,
+ },
+ }).get();
+ }
+
+ async tool(iteration: BeeIterationToolResult): Promise<{ output: string; success: boolean }> {
+ const tool = this.input.tools.find(
+ (tool) => tool.name.trim().toUpperCase() == iteration.tool_name?.toUpperCase(),
+ );
+ if (!tool) {
+ this.failedAttemptsCounter.use();
+ const availableTools = this.input.tools.map((tool) => tool.name);
+ return {
+ success: false,
+ output: [
+ `Tool does not exist!`,
+ availableTools.length > 0 &&
+ `Use one of the following tools: ${availableTools.join(",")}`,
+ ]
+ .filter(Boolean)
+ .join("\n"),
+ };
+ }
+ const options = await (async () => {
+ const baseOptions: BaseToolRunOptions = {
+ signal: this.run.signal,
+ };
+ const customOptions = await this.options.modifiers?.getToolRunOptions?.({
+ tool,
+ input: iteration.tool_input,
+ baseOptions,
+ });
+ return customOptions ?? baseOptions;
+ })();
+
+ return new Retryable({
+ config: {
+ signal: this.run.signal,
+ maxRetries: this.options.execution?.maxRetriesPerStep,
+ },
+ onError: async (error) => {
+ await this.run.emitter.emit("toolError", {
+ data: {
+ iteration,
+ tool,
+ input: iteration.tool_input,
+ options,
+ error: FrameworkError.ensure(error),
+ },
+ });
+ this.failedAttemptsCounter.use(error);
+ },
+ executor: async () => {
+ await this.run.emitter.emit("toolStart", {
+ data: {
+ tool,
+ input: iteration.tool_input,
+ options,
+ iteration,
+ },
+ });
+
+ try {
+ const toolOutput: ToolOutput = await tool.run(iteration.tool_input, options);
+ await this.run.emitter.emit("toolSuccess", {
+ data: {
+ tool,
+ input: iteration.tool_input,
+ options,
+ result: toolOutput,
+ iteration,
+ },
+ });
+
+ if (toolOutput.isEmpty()) {
+ return { output: BeeToolNoResultsPrompt.render({}), success: true };
+ }
+
+ return {
+ success: true,
+ output: toolOutput.getTextContent(),
+ };
+ } catch (error) {
+ if (error instanceof ToolInputValidationError) {
+ this.failedAttemptsCounter.use(error);
+ return {
+ success: false,
+ output: BeeToolInputErrorPrompt.render({
+ reason: error.toString(),
+ }),
+ };
+ }
+
+ await this.run.emitter.emit("toolError", {
+ data: {
+ tool,
+ input: iteration.tool_input,
+ options,
+ error,
+ iteration,
+ },
+ });
+
+ if (FrameworkError.isRetryable(error)) {
+ this.failedAttemptsCounter.use(error);
+ return {
+ success: false,
+ output: BeeToolErrorPrompt.render({
+ reason: FrameworkError.ensure(error).explain(),
+ }),
+ };
+ }
+
+ throw error;
+ }
+ },
+ }).get();
+ }
+}
diff --git a/src/agents/bee/types.ts b/src/agents/bee/types.ts
new file mode 100644
index 00000000..46b8e8c0
--- /dev/null
+++ b/src/agents/bee/types.ts
@@ -0,0 +1,102 @@
+/**
+ * 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 { ChatLLMOutput } from "@/llms/chat.js";
+import { BeeIterationResult, BeeIterationToolResult } from "@/agents/bee/parser.js";
+import { BaseMemory } from "@/memory/base.js";
+import { BaseMessage } from "@/llms/primitives/message.js";
+import { Callback } from "@/emitter/types.js";
+import { AnyTool, BaseToolRunOptions, Tool, ToolError, ToolOutput } from "@/tools/base.js";
+
+export interface BeeRunInput {
+ prompt: string;
+}
+
+export interface BeeRunOutput {
+ result: BaseMessage;
+ iterations: BeeAgentRunIteration[];
+ memory: BaseMemory;
+}
+
+export interface BeeAgentRunIteration {
+ raw: ChatLLMOutput;
+ state: BeeIterationResult;
+}
+
+export interface BeeAgentExecutionConfig {
+ totalMaxRetries?: number;
+ maxRetriesPerStep?: number;
+ maxIterations?: number;
+}
+
+export interface BeeRunOptions {
+ signal?: AbortSignal;
+ execution?: BeeAgentExecutionConfig;
+ modifiers?: {
+ getToolRunOptions?: (execution: {
+ tool: Tool ;
+ input: unknown;
+ baseOptions: BaseToolRunOptions;
+ }) => BaseToolRunOptions | Promise;
+ };
+}
+
+export interface BeeUpdateMeta {
+ success: boolean;
+}
+
+export interface BeeCallbacks {
+ start?: Callback;
+ error?: Callback<{ error: Error }>;
+ retry?: Callback;
+ success?: Callback<{ data: BaseMessage }>;
+ update?: Callback<{
+ data: BeeIterationResult;
+ update: { key: keyof BeeIterationResult; value: string };
+ meta: BeeUpdateMeta;
+ }>;
+ partialUpdate?: Callback<{
+ data: BeeIterationResult;
+ update: { key: keyof BeeIterationResult; value: string };
+ meta: BeeUpdateMeta;
+ }>;
+ toolStart?: Callback<{
+ data: {
+ tool: AnyTool;
+ input: unknown;
+ options: BaseToolRunOptions;
+ iteration: BeeIterationToolResult;
+ };
+ }>;
+ toolSuccess?: Callback<{
+ data: {
+ tool: AnyTool;
+ input: unknown;
+ options: BaseToolRunOptions;
+ result: ToolOutput;
+ iteration: BeeIterationToolResult;
+ };
+ }>;
+ toolError?: Callback<{
+ data: {
+ tool: AnyTool;
+ input: unknown;
+ options: BaseToolRunOptions;
+ error: ToolError;
+ iteration: BeeIterationToolResult;
+ };
+ }>;
+}
diff --git a/src/agents/manager.test.ts b/src/agents/manager.test.ts
new file mode 100644
index 00000000..7112ef87
--- /dev/null
+++ b/src/agents/manager.test.ts
@@ -0,0 +1,53 @@
+/**
+ * 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 { BaseAgentManager } from "@/agents/manager.js";
+import { expect } from "vitest";
+
+describe("Agent Manager", () => {
+ class AgentManager extends BaseAgentManager {
+ protected _create(overrides: any) {
+ return overrides;
+ }
+ }
+
+ it("Creates", () => {
+ const manager = new AgentManager();
+ const agents = [{ a: "a" }, { b: "b" }, { c: "c" }];
+ for (const agentInput of agents) {
+ const agent = manager.create(agentInput);
+ expect(agent).toBe(agentInput);
+ }
+ });
+
+ it("Destroys", () => {
+ const manager = new AgentManager();
+
+ const agents = [
+ manager.create({
+ destroy: vi.fn(),
+ }),
+ manager.create({
+ destroy: vi.fn(),
+ }),
+ ];
+
+ manager.destroy();
+ for (const agent of agents) {
+ expect(agent.destroy).toBeCalled();
+ }
+ });
+});
diff --git a/src/agents/manager.ts b/src/agents/manager.ts
new file mode 100644
index 00000000..ab1aa96d
--- /dev/null
+++ b/src/agents/manager.ts
@@ -0,0 +1,51 @@
+/**
+ * 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 { BaseAgent } from "@/agents/base.js";
+
+export abstract class BaseAgentManager<
+ T extends BaseAgent,
+ PDefault extends readonly any[],
+ POverride,
+> {
+ public readonly instances: Set>;
+ public readonly defaults: PDefault;
+
+ constructor(...defaults: PDefault) {
+ this.defaults = defaults;
+ this.instances = new Set>();
+ }
+
+ protected abstract _create(overrides: POverride): T;
+
+ create(overrides: POverride): T {
+ const instance = this._create(overrides);
+ const ref = new WeakRef(instance);
+ this.instances.add(ref);
+ return instance;
+ }
+
+ destroy() {
+ for (const weakRef of Array.from(this.instances.values())) {
+ this.instances.delete(weakRef);
+
+ const instance = weakRef.deref();
+ if (instance) {
+ instance.destroy();
+ }
+ }
+ }
+}
diff --git a/src/agents/types.ts b/src/agents/types.ts
new file mode 100644
index 00000000..571ef4e2
--- /dev/null
+++ b/src/agents/types.ts
@@ -0,0 +1,43 @@
+/**
+ * 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 { BaseAgent } from "@/agents/base.js";
+import { OmitEmpty } from "@/internals/types.js";
+
+export interface AgentMeta {
+ name: string;
+ description: string;
+ extraDescription?: string;
+}
+
+export type AgentCallbackValue =
+ | { data?: never; error: Error }
+ | { data: unknown; error?: never }
+ | object;
+
+export type InternalAgentCallbackValue<
+ T extends AgentCallbackValue,
+ E extends NonNullable,
+> = OmitEmpty & E;
+
+export type PublicAgentCallbackValue =
+ OmitEmpty;
+
+export type AgentCallback = (value: T) => void;
+
+export type GetAgentInput = T extends BaseAgent ? X : never;
+export type GetAgentOutput = T extends BaseAgent ? X : never;
+export type AnyAgent = BaseAgent;
diff --git a/src/cache/base.ts b/src/cache/base.ts
new file mode 100644
index 00000000..45d98118
--- /dev/null
+++ b/src/cache/base.ts
@@ -0,0 +1,28 @@
+/**
+ * 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 { Serializable } from "@/internals/serializable.js";
+
+export abstract class BaseCache extends Serializable {
+ public enabled = true;
+
+ abstract size(): Promise;
+ abstract set(key: string, value: T): Promise;
+ abstract get(key: string): Promise;
+ abstract has(key: string): Promise;
+ abstract delete(key: string): Promise;
+ abstract clear(): Promise;
+}
diff --git a/src/cache/decoratorCache.test.ts b/src/cache/decoratorCache.test.ts
new file mode 100644
index 00000000..82b9a363
--- /dev/null
+++ b/src/cache/decoratorCache.test.ts
@@ -0,0 +1,285 @@
+/**
+ * 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 {
+ Cache,
+ SingletonCacheKeyFn,
+ JSONCacheKeyFn,
+ ObjectHashKeyFn,
+ CacheFn,
+} from "@/cache/decoratorCache.js";
+import type { AnyFn } from "@/internals/types.js";
+
+describe("@Cache decorator", () => {
+ it("Caches method", () => {
+ class A {
+ static number = 0;
+
+ @Cache()
+ getNumber() {
+ A.number += 1;
+ return A.number;
+ }
+ }
+
+ const a = new A();
+ expect(a.getNumber()).toEqual(1);
+ expect(a.getNumber()).toEqual(1);
+ });
+
+ it("Caches a complex datatype", () => {
+ class A {
+ static number = 0;
+
+ @Cache({
+ cacheKey: ObjectHashKeyFn,
+ })
+ getNumber(handler: AnyFn) {
+ A.number += 1;
+ return handler;
+ }
+ }
+
+ const a = new A();
+ const x = () => 1;
+ const y = () => 1;
+
+ expect(a.getNumber(x)).toEqual(x);
+ expect(A.number).toBe(1);
+ expect(a.getNumber(y)).toEqual(y);
+ expect(A.number).toBe(2);
+ expect(a.getNumber(x)).toEqual(x);
+ expect(A.number).toBe(2);
+ expect(a.getNumber(y)).toEqual(y);
+ expect(A.number).toBe(2);
+ });
+
+ it("Ignores args if 'EmptyCacheKeyFn' is specified.", () => {
+ class A {
+ static number = 0;
+
+ @Cache({
+ cacheKey: SingletonCacheKeyFn,
+ })
+ // eslint-disable-next-line unused-imports/no-unused-vars
+ getNumber(...args: any[]) {
+ A.number += 1;
+ return A.number;
+ }
+ }
+
+ const a = new A();
+ expect(a.getNumber("a", "b")).toEqual(1);
+ expect(a.getNumber("b", "c")).toEqual(1);
+ });
+
+ it("Caches an accessor", () => {
+ class A {
+ static number = 0;
+
+ @Cache()
+ get getNumber() {
+ A.number += 1;
+ return A.number;
+ }
+ }
+
+ const a = new A();
+ expect(a.getNumber).toEqual(1);
+ expect(a.getNumber).toEqual(1);
+ });
+
+ it("Controls cache behaviour", () => {
+ class A {
+ constructor(public invocations = 0) {}
+
+ @Cache()
+ get getNumber() {
+ this.invocations += 1;
+ return this.invocations;
+ }
+ }
+
+ const a = new A();
+ const cache = Cache.getInstance(a, "getNumber");
+ cache.disable();
+ expect(a.getNumber).toEqual(1);
+ expect(a.getNumber).toEqual(2);
+ expect(a.getNumber).toEqual(3);
+ cache.enable();
+ expect(a.getNumber).toEqual(4);
+ expect(a.getNumber).toEqual(4);
+ });
+
+ it("Clears cache", () => {
+ class A {
+ static number = 0;
+
+ @Cache()
+ get getNumber() {
+ A.number += 1;
+ return A.number;
+ }
+
+ reset() {
+ Cache.getInstance(this, "getNumber").clear();
+ }
+ }
+
+ const a = new A();
+ expect(a.getNumber).toEqual(1);
+ expect(a.getNumber).toEqual(1);
+ a.reset();
+ expect(a.getNumber).toEqual(2);
+ });
+
+ it("Clears cache by a key", () => {
+ class A {
+ static calls = 0;
+
+ @Cache({
+ cacheKey: JSONCacheKeyFn,
+ })
+ getNumber(a: number) {
+ A.calls += a;
+ return A.calls;
+ }
+
+ reset(values: number[]) {
+ Cache.getInstance(this, "getNumber").clear(values.map((value) => JSONCacheKeyFn(value)));
+ }
+ }
+
+ const a = new A();
+ expect(a.getNumber(10)).toEqual(10);
+ expect(a.getNumber(10)).toEqual(10);
+ expect(a.getNumber(20)).toEqual(30);
+ expect(a.getNumber(20)).toEqual(30);
+ a.reset([10]);
+ expect(a.getNumber(10)).toEqual(40);
+ expect(a.getNumber(20)).toEqual(30);
+ });
+
+ it("Caches a complex input", () => {
+ class A {
+ static calls = 0;
+
+ @Cache()
+ calculate(a: number, b: number) {
+ A.calls += 1;
+ return a + b;
+ }
+ }
+
+ const a = new A();
+ expect(a.calculate(2, 3)).toEqual(5);
+ expect(a.calculate(2, 3)).toEqual(5);
+ expect(A.calls).toEqual(1);
+ });
+
+ it("Does not interfere cache across instances", () => {
+ class A {
+ constructor(public readonly base: number) {}
+
+ @Cache()
+ add(input: number) {
+ return input + this.base;
+ }
+ }
+
+ const a = new A(100);
+ const b = new A(200);
+
+ expect(a.add(1)).toEqual(101);
+ expect(a.add(2)).toEqual(102);
+
+ expect(b.add(1)).toEqual(201);
+ expect(b.add(2)).toEqual(202);
+ });
+
+ it("Preserves meta information", () => {
+ class A {
+ @Cache()
+ calculate() {}
+ }
+
+ const a = new A();
+ const b = new A();
+ expect(a.calculate.name).toBe("calculate");
+ expect(b.calculate.name).toBe("calculate");
+ });
+
+ it("Static caching", () => {
+ class A {
+ static counter = 0;
+
+ @Cache()
+ static add(input: number) {
+ A.counter += 1;
+ return input;
+ }
+
+ static bbb(input: number) {
+ A.counter += 1;
+ return input;
+ }
+ }
+
+ class B extends A {}
+
+ expect(A.add(1)).toEqual(1);
+ expect(A.add(1)).toEqual(1);
+ expect(B.add(1)).toEqual(1);
+ expect(B.add(1)).toEqual(1);
+
+ expect(A.counter).toBe(2);
+ expect(B.counter).toBe(2);
+ });
+
+ describe("CacheFn", () => {
+ it("Caches", () => {
+ const fn = vi.fn((input) => input);
+
+ const cached = CacheFn.create(fn);
+ expect(cached(1)).toBe(1);
+ expect(cached(1)).toBe(1);
+ expect(cached(2)).toBe(2);
+ expect(cached(2)).toBe(2);
+ expect(cached(3)).toBe(3);
+ expect(cached(3)).toBe(3);
+
+ expect(fn).toHaveBeenCalledTimes(3);
+ });
+
+ it("Updates TTL", async () => {
+ const sleep = async (ms: number) => await new Promise((resolve) => setTimeout(resolve, ms));
+ let counter = 0;
+ const cached = CacheFn.create(
+ vi.fn(async () => {
+ await sleep(50);
+ cached.updateTTL(200);
+ counter += 1;
+ return counter;
+ }),
+ );
+ await expect(cached()).resolves.toBe(1);
+ await expect(cached()).resolves.toBe(1);
+ await sleep(250);
+ await expect(cached()).resolves.toBe(2);
+ await expect(cached()).resolves.toBe(2);
+ });
+ });
+});
diff --git a/src/cache/decoratorCache.ts b/src/cache/decoratorCache.ts
new file mode 100644
index 00000000..06c0429b
--- /dev/null
+++ b/src/cache/decoratorCache.ts
@@ -0,0 +1,302 @@
+/**
+ * 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 { AnyFn, TypedFn } from "@/internals/types.js";
+import * as R from "remeda";
+import hash from "object-hash";
+import { createHash, createRandomHash } from "@/internals/helpers/hash.js";
+
+type InputFn = AnyFn;
+type TargetFn = AnyFn;
+type Instance = NonNullable;
+
+type CacheKeyFn = (...args: any[]) => string;
+
+export interface CacheDecoratorOptions {
+ enabled: boolean;
+ ttl?: number;
+ cacheKey: CacheKeyFn;
+ enumerable?: boolean;
+}
+
+interface FnContext {
+ options: CacheDecoratorOptions;
+ cache: Map;
+}
+
+interface GroupContext {
+ options: CacheDecoratorOptions;
+ instances: WeakMap;
+}
+
+export interface CacheDecoratorInstance {
+ clear(keys?: string[]): void;
+
+ isEnabled(): boolean;
+
+ enable(): void;
+
+ disable(): void;
+
+ update(data: Partial): void;
+}
+
+const state = {
+ container: new WeakMap(),
+ extractDescriptor(descriptor: PropertyDescriptor) {
+ // Method approach
+ if (descriptor.value != null) {
+ return {
+ value: descriptor.value as AnyFn,
+ update: (value: TargetFn) => {
+ descriptor.value = value;
+ },
+ };
+ }
+ // Accessor approach
+ if (descriptor.get != null) {
+ return {
+ value: descriptor.get as AnyFn,
+ update: (value: TargetFn) => {
+ descriptor.get = value;
+ },
+ };
+ }
+ throw new Error(`@${Cache.name} decorator must be either on a method or get accessor.`);
+ },
+ getInstanceContext>(target: T, ctx: GroupContext) {
+ if (!ctx.instances.has(target)) {
+ ctx.instances.set(target, {
+ cache: new Map(),
+ options: ctx.options,
+ });
+ }
+ return ctx.instances.get(target)!;
+ },
+} as const;
+
+const initQueue = new WeakMap();
+
+export function Cache(_options?: Partial) {
+ const baseOptions: Required = {
+ enabled: true,
+ cacheKey: ObjectHashKeyFn,
+ ttl: Infinity,
+ enumerable: false,
+ ...R.pickBy(_options ?? {}, R.isDefined),
+ };
+
+ return function handler(
+ obj: NonNullable,
+ key: string | symbol,
+ descriptor: PropertyDescriptor,
+ ) {
+ if (Object.hasOwn(obj, "constructor")) {
+ const constructor = obj.constructor;
+ if (!initQueue.has(constructor)) {
+ initQueue.set(constructor, []);
+ }
+
+ baseOptions.enumerable = Boolean(descriptor.get);
+ initQueue
+ .get(constructor)!
+ .push({ key, enumerable: _options?.enumerable ?? baseOptions.enumerable });
+ }
+
+ const target = state.extractDescriptor(descriptor);
+ const groupContext: GroupContext = { instances: new WeakMap(), options: baseOptions };
+
+ const fn = function wrapper(this: any, ...args: any[]) {
+ const invokeOriginal = () => target.value.apply(this, args);
+ const ctx = state.getInstanceContext(this, groupContext);
+ if (!ctx.options.enabled) {
+ return invokeOriginal();
+ }
+
+ const inputHash = ctx.options.cacheKey(...args);
+ if (
+ !ctx.cache.has(inputHash) ||
+ (ctx.cache.get(inputHash)?.expiresAt ?? Infinity) < Date.now() // is expired check
+ ) {
+ const result = invokeOriginal();
+ ctx.cache.set(inputHash, {
+ expiresAt: Date.now() + (ctx.options.ttl ?? Infinity),
+ data: result,
+ });
+ }
+ return ctx.cache.get(inputHash)!.data;
+ };
+ Object.defineProperty(fn, "name", {
+ get: () => target.value.name ?? "anonymous",
+ });
+
+ target.update(fn);
+ state.container.set(fn, groupContext);
+ };
+}
+
+Cache.init = function init>(self: T) {
+ const task = initQueue.get(self.constructor) ?? [];
+ for (const { key, enumerable } of task) {
+ const descriptor = Object.getOwnPropertyDescriptor(self.constructor.prototype, key);
+ if (descriptor) {
+ Object.defineProperty(self, key, Object.assign(descriptor, { enumerable }));
+ }
+ }
+ initQueue.delete(self);
+};
+
+Cache.getInstance = function getInstance>(
+ target: T,
+ property: keyof T,
+): CacheDecoratorInstance {
+ const descriptor = Object.getOwnPropertyDescriptor(Object.getPrototypeOf(target), property);
+ if (!descriptor) {
+ throw new TypeError(`No descriptor has been found for '${String(property)}'`);
+ }
+ const value = state.extractDescriptor(descriptor);
+ const ctxByInstance = state.container.get(value.value);
+ if (!ctxByInstance) {
+ throw new TypeError(`No cache instance is bounded to '${String(property)}'!`);
+ }
+
+ const ctx = state.getInstanceContext(target, ctxByInstance);
+ return {
+ clear(keys?: string[]) {
+ if (keys) {
+ keys.forEach((key) => ctx.cache.delete(key));
+ } else {
+ ctx.cache.clear();
+ }
+ },
+ update(data: CacheDecoratorOptions) {
+ const oldTTL = ctx.options.ttl;
+ const newTTL = data.ttl;
+ if (oldTTL !== newTTL && newTTL !== undefined) {
+ for (const value of ctx.cache.values()) {
+ if (value.expiresAt === Infinity) {
+ value.expiresAt = Date.now() + newTTL;
+ } else {
+ const diff = newTTL - (oldTTL ?? 0);
+ value.expiresAt += diff;
+ }
+ }
+ }
+ Object.assign(ctx.options, data);
+ },
+ isEnabled() {
+ return ctx.options.enabled;
+ },
+ enable() {
+ ctx.options.enabled = true;
+ },
+ disable() {
+ ctx.options.enabled = false;
+ },
+ };
+};
+
+export const WeakRefKeyFn: CacheKeyFn = (() => {
+ const _lookup = new WeakMap();
+
+ return (...args: any[]) => {
+ const chunks = args.map((value) => {
+ if (R.isObjectType(value) || R.isFunction(value)) {
+ if (!_lookup.has(value)) {
+ _lookup.set(value, createRandomHash(6));
+ }
+ return _lookup.get(value)!;
+ }
+ return value;
+ });
+ return createHash(JSON.stringify(chunks));
+ };
+})();
+
+export const ObjectHashKeyFn: CacheKeyFn = (...args: any[]) =>
+ hash(args, {
+ encoding: "base64",
+ replacer: (() => {
+ const _lookup = new WeakMap();
+ return (value: unknown) => {
+ if (value && value instanceof AbortSignal) {
+ // not supported by "hash" function
+ if (!_lookup.has(value)) {
+ _lookup.set(value, createRandomHash(6));
+ }
+ return _lookup.get(value)!;
+ }
+ return value;
+ };
+ })(),
+ unorderedArrays: false,
+ unorderedObjects: false,
+ unorderedSets: false,
+ });
+export const JSONCacheKeyFn: CacheKeyFn = (...args: any[]) => JSON.stringify(args);
+// eslint-disable-next-line unused-imports/no-unused-vars
+export const SingletonCacheKeyFn: CacheKeyFn = (...args: any[]) => "";
+
+export class CacheFn extends Function {
+ readonly name = CacheFn.name;
+
+ static create(
+ fn: (...args: A) => B,
+ options?: Partial,
+ ) {
+ const instance = new CacheFn(fn, options);
+ return instance as TypedFn & CacheFn ;
+ }
+
+ constructor(
+ protected readonly fn: (...args: P) => R,
+ protected readonly options?: Partial,
+ ) {
+ super();
+
+ Cache.getInstance(this, "get").update(options ?? {});
+ return new Proxy(this, {
+ apply(this: CacheFn, target, _, args: any[]) {
+ return target.get(...(args as unknown as P));
+ },
+ get(target, prop: string | symbol, receiver: any) {
+ const value = target[prop as keyof typeof target];
+ if (value instanceof Function) {
+ return function (this: CacheFn
, ...args: P) {
+ return value.apply(this === receiver ? target : this, args);
+ };
+ }
+ return value;
+ },
+ }) as unknown as CacheFn
& typeof fn;
+ }
+
+ updateTTL(ttl: number) {
+ Cache.getInstance(this, "get").update({ ttl });
+ }
+
+ createSnapshot() {
+ return {
+ fn: this.fn,
+ options: this.options,
+ };
+ }
+
+ @Cache()
+ get(...args: P) {
+ return this.fn(...args);
+ }
+}
diff --git a/src/cache/fileCache.ts b/src/cache/fileCache.ts
new file mode 100644
index 00000000..0fcf3a4c
--- /dev/null
+++ b/src/cache/fileCache.ts
@@ -0,0 +1,123 @@
+/**
+ * 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 { BaseCache } from "@/cache/base.js";
+import fs from "node:fs";
+import { SlidingCache } from "@/cache/slidingCache.js";
+import { Cache } from "@/cache/decoratorCache.js";
+import { Serializer } from "@/serializer/serializer.js";
+import { SerializableClass } from "@/internals/serializable.js";
+
+interface Input {
+ fullPath: string;
+}
+
+export class FileCache extends BaseCache {
+ constructor(protected readonly input: Input) {
+ super();
+ }
+
+ static {
+ this.register();
+ }
+
+ static async fromProvider(provider: BaseCache , input: Input) {
+ await fs.promises.writeFile(input.fullPath, provider.serialize());
+ return new FileCache(input);
+ }
+
+ @Cache()
+ protected async getProvider(): Promise> {
+ const exists = await fs.promises
+ .stat(this.input.fullPath)
+ .then((r) => r.isFile())
+ .catch(() => false);
+
+ if (exists) {
+ const serialized = await fs.promises.readFile(this.input.fullPath, "utf8");
+ const { target, snapshot } = await Serializer.deserialize(serialized);
+ const Target = Serializer.getFactory(target).ref as SerializableClass>;
+ const instance = Target.fromSnapshot(snapshot);
+ if (!(instance instanceof BaseCache)) {
+ throw new TypeError("Provided file does not serialize any instance of BaseCache class.");
+ }
+ return instance;
+ } else {
+ return new SlidingCache({
+ size: Infinity,
+ ttl: Infinity,
+ });
+ }
+ }
+
+ async reload() {
+ // @ts-expect-error protected property
+ Cache.getInstance(this, "getProvider").clear();
+ void (await this.getProvider());
+ }
+
+ protected async save() {
+ const provider = await this.getProvider();
+ return await fs.promises.writeFile(this.input.fullPath, provider.serialize());
+ }
+
+ async size() {
+ const provider = await this.getProvider();
+ return provider.size();
+ }
+
+ async set(key: string, value: T) {
+ const provider = await this.getProvider();
+ await provider.set(key, value);
+ void provider.get(key).finally(() => {
+ void this.save();
+ });
+ }
+
+ async get(key: string) {
+ const provider = await this.getProvider();
+ return await provider.get(key);
+ }
+
+ async has(key: string) {
+ const provider = await this.getProvider();
+ return await provider.has(key);
+ }
+
+ async delete(key: string) {
+ const provider = await this.getProvider();
+ const result = await provider.delete(key);
+ await this.save();
+ return result;
+ }
+
+ async clear() {
+ const provider = await this.getProvider();
+ await provider.clear();
+ await this.save();
+ }
+
+ async createSnapshot() {
+ return {
+ input: { fullPath: this.input.fullPath },
+ provider: await this.getProvider(),
+ };
+ }
+
+ loadSnapshot(snapshot: Awaited>): void {
+ Object.assign(this, snapshot);
+ }
+}
diff --git a/src/cache/nullCache.test.ts b/src/cache/nullCache.test.ts
new file mode 100644
index 00000000..ad39bff0
--- /dev/null
+++ b/src/cache/nullCache.test.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 { NullCache } from "@/cache/nullCache.js";
+import { verifyDeserialization } from "@tests/e2e/utils.js";
+
+describe("NullCache", () => {
+ it("Handles basic operations", async () => {
+ const instance = new NullCache();
+ await instance.set("1", 1);
+ await expect(instance.has("1")).resolves.toBe(false);
+ await expect(instance.get("1")).resolves.toBe(undefined);
+ await instance.delete("1");
+ await expect(instance.size()).resolves.toBe(0);
+ await instance.clear();
+ });
+
+ it("Serializes", async () => {
+ const instance = new NullCache();
+ await instance.set("1", 1);
+ const serialized = instance.serialize();
+ const deserialized = NullCache.fromSerialized(serialized);
+ verifyDeserialization(instance, deserialized);
+ });
+});
diff --git a/src/cache/nullCache.ts b/src/cache/nullCache.ts
new file mode 100644
index 00000000..1bf7dfed
--- /dev/null
+++ b/src/cache/nullCache.ts
@@ -0,0 +1,51 @@
+/**
+ * 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 { BaseCache } from "@/cache/base.js";
+
+export class NullCache extends BaseCache {
+ enabled = false;
+
+ async set(_key: string, _value: T) {}
+
+ async get(_key: string) {
+ return undefined;
+ }
+
+ async has(_key: string) {
+ return false;
+ }
+
+ async delete(_key: string) {
+ return true;
+ }
+
+ async clear() {}
+
+ async size() {
+ return 0;
+ }
+
+ createSnapshot() {
+ return {
+ enabled: this.enabled,
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
diff --git a/src/cache/slidingCache.test.ts b/src/cache/slidingCache.test.ts
new file mode 100644
index 00000000..f288a3f7
--- /dev/null
+++ b/src/cache/slidingCache.test.ts
@@ -0,0 +1,73 @@
+/**
+ * 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 { verifyDeserialization } from "@tests/e2e/utils.js";
+import { SlidingCache } from "@/cache/slidingCache.js";
+import { beforeEach } from "vitest";
+
+describe("SlidingCache", () => {
+ beforeEach(() => {
+ vitest.useFakeTimers();
+ });
+
+ it("Handles basic operations", async () => {
+ const instance = new SlidingCache({
+ ttl: 5000,
+ size: 10,
+ });
+ await instance.set("1", 1);
+ await instance.set("2", 2);
+ await expect(instance.has("1")).resolves.toBe(true);
+ await expect(instance.get("1")).resolves.toBe(1);
+ await expect(instance.size()).resolves.toBe(2);
+ await instance.delete("1");
+ await expect(instance.size()).resolves.toBe(1);
+ await instance.clear();
+ await expect(instance.size()).resolves.toBe(0);
+ });
+
+ it("Removes old entries", async () => {
+ const instance = new SlidingCache({
+ ttl: 2500,
+ size: 10,
+ });
+
+ await instance.set("1", 1);
+ await vitest.advanceTimersByTimeAsync(1500);
+ await expect(instance.size()).resolves.toBe(1);
+ await expect(instance.has("1")).resolves.toBe(true);
+ await expect(instance.get("1")).resolves.toBe(1);
+
+ await vitest.advanceTimersByTimeAsync(2000);
+ await expect(instance.has("1")).resolves.toBe(false);
+ await expect(instance.get("1")).resolves.toBe(undefined);
+
+ await instance.clear();
+ });
+
+ it("Serializes", async () => {
+ const instance = new SlidingCache({
+ ttl: 5000,
+ size: 10,
+ });
+ await instance.set("1", 1);
+ await instance.set("2", 2);
+ await instance.set("3", 3);
+ const serialized = instance.serialize();
+ const deserialized = SlidingCache.fromSerialized(serialized);
+ verifyDeserialization(instance, deserialized);
+ });
+});
diff --git a/src/cache/slidingCache.ts b/src/cache/slidingCache.ts
new file mode 100644
index 00000000..49c5cf72
--- /dev/null
+++ b/src/cache/slidingCache.ts
@@ -0,0 +1,101 @@
+/**
+ * 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 { SlidingTaskMap, Task } from "promise-based-task";
+import { BaseCache } from "@/cache/base.js";
+import { Serializable } from "@/internals/serializable.js";
+
+export interface SlidingCacheInput {
+ size: number;
+ ttl?: number;
+}
+
+class SlidingCacheEntry extends Serializable {
+ constructor(protected readonly value: T) {
+ super();
+ }
+
+ static {
+ this.register();
+ }
+
+ destructor() {
+ if (this.value instanceof Task) {
+ this.value.destructor();
+ }
+ }
+
+ unwrap(): T {
+ return this.value;
+ }
+
+ createSnapshot() {
+ return { value: this.value };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
+
+export class SlidingCache extends BaseCache {
+ protected readonly provider;
+
+ constructor(input: SlidingCacheInput) {
+ super();
+ this.provider = new SlidingTaskMap>(input.size, input.ttl);
+ }
+
+ static {
+ this.register();
+ }
+
+ async get(key: string) {
+ const value = this.provider.get(key);
+ return value?.unwrap?.();
+ }
+
+ async has(key: string) {
+ return this.provider.has(key);
+ }
+
+ async clear() {
+ this.provider.clear();
+ }
+
+ async delete(key: string) {
+ return this.provider.delete(key);
+ }
+
+ async set(key: string, value: T) {
+ this.provider.set(key, new SlidingCacheEntry(value));
+ }
+
+ async size() {
+ return this.provider.size;
+ }
+
+ createSnapshot() {
+ return {
+ enabled: this.enabled,
+ provider: this.provider,
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
diff --git a/src/cache/unconstrainedCache.test.ts b/src/cache/unconstrainedCache.test.ts
new file mode 100644
index 00000000..09fd899c
--- /dev/null
+++ b/src/cache/unconstrainedCache.test.ts
@@ -0,0 +1,46 @@
+/**
+ * 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 { verifyDeserialization } from "@tests/e2e/utils.js";
+import { beforeEach } from "vitest";
+import { UnconstrainedCache } from "@/cache/unconstrainedCache.js";
+
+describe("UnconstrainedCache", () => {
+ beforeEach(() => {
+ vitest.useFakeTimers();
+ });
+
+ it("Handles basic operations", async () => {
+ const instance = new UnconstrainedCache();
+ await instance.set("1", 1);
+ await instance.set("2", 2);
+ await expect(instance.has("1")).resolves.toBe(true);
+ await expect(instance.get("1")).resolves.toBe(1);
+ await expect(instance.size()).resolves.toBe(2);
+ await instance.delete("1");
+ await expect(instance.size()).resolves.toBe(1);
+ await instance.clear();
+ await expect(instance.size()).resolves.toBe(0);
+ });
+
+ it("Serializes", async () => {
+ const instance = new UnconstrainedCache();
+ await instance.set("1", 1);
+ const serialized = instance.serialize();
+ const deserialized = UnconstrainedCache.fromSerialized(serialized);
+ verifyDeserialization(instance, deserialized);
+ });
+});
diff --git a/src/cache/unconstrainedCache.ts b/src/cache/unconstrainedCache.ts
new file mode 100644
index 00000000..19589d38
--- /dev/null
+++ b/src/cache/unconstrainedCache.ts
@@ -0,0 +1,56 @@
+/**
+ * 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 { BaseCache } from "@/cache/base.js";
+
+export class UnconstrainedCache extends BaseCache {
+ protected readonly provider = new Map();
+
+ async get(key: string) {
+ return this.provider.get(key);
+ }
+
+ async has(key: string) {
+ return this.provider.has(key);
+ }
+
+ async clear() {
+ this.provider.clear();
+ }
+
+ async delete(key: string) {
+ return this.provider.delete(key);
+ }
+
+ async set(key: string, value: T) {
+ this.provider.set(key, value);
+ }
+
+ async size() {
+ return this.provider.size;
+ }
+
+ createSnapshot() {
+ return {
+ enabled: this.enabled,
+ provider: this.provider,
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
diff --git a/src/context.ts b/src/context.ts
new file mode 100644
index 00000000..0e693c84
--- /dev/null
+++ b/src/context.ts
@@ -0,0 +1,175 @@
+/**
+ * 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 { AsyncLocalStorage } from "node:async_hooks";
+import { Emitter } from "@/emitter/emitter.js";
+import { createRandomHash } from "@/internals/helpers/hash.js";
+import { omit } from "remeda";
+import { Callback } from "@/emitter/types.js";
+import { createNonOverridableObject } from "@/internals/helpers/object.js";
+import { registerSignals } from "@/internals/helpers/cancellation.js";
+import { Serializable } from "@/internals/serializable.js";
+import { LazyPromise } from "@/internals/helpers/promise.js";
+import { FrameworkError } from "@/errors.js";
+
+export interface RunInstance {
+ emitter: Emitter;
+}
+
+export interface RunContextCallbacks {
+ start: Callback;
+ success: Callback;
+ error: Callback;
+ finish: Callback;
+}
+
+export type GetRunContext = T extends RunInstance ? RunContext : never;
+export type GetRunInstance = T extends RunInstance ? P : never;
+
+export class Run extends LazyPromise {
+ constructor(
+ handler: () => Promise,
+ protected readonly runContext: RunContext,
+ ) {
+ super(handler);
+ }
+
+ readonly [Symbol.toStringTag] = "Promise";
+
+ observe(fn: (emitter: Emitter) => void) {
+ fn(this.runContext.emitter as any);
+ return this;
+ }
+
+ context(value: object) {
+ Object.assign(this.runContext.context, value);
+ Object.assign(this.runContext.emitter.context, value);
+ return this;
+ }
+
+ middleware(fn: (context: RunContext) => void) {
+ fn(this.runContext);
+ return this;
+ }
+}
+
+export class RunContext extends Serializable {
+ static #storage = new AsyncLocalStorage>();
+
+ protected readonly controller: AbortController;
+ public readonly runId: string;
+ public readonly groupId: string;
+ public readonly parentId?: string;
+ public readonly emitter;
+ public readonly context: object;
+
+ get signal() {
+ return this.controller.signal;
+ }
+
+ abort() {
+ this.controller.abort();
+ }
+
+ constructor(
+ protected readonly instance: RunInstance,
+ parent?: RunContext,
+ signal?: AbortSignal,
+ ) {
+ super();
+ this.runId = createRandomHash(5);
+ this.parentId = parent?.runId;
+ this.groupId = parent?.groupId ?? createRandomHash();
+ this.context = createNonOverridableObject(
+ omit((parent?.context ?? {}) as any, ["id", "parentId"]),
+ );
+
+ this.controller = new AbortController();
+ registerSignals(this.controller, [signal, parent?.signal]);
+
+ this.emitter = instance.emitter.child({
+ context: this.context,
+ trace: {
+ id: this.groupId,
+ runId: this.runId,
+ parentRunId: parent?.runId,
+ },
+ });
+ if (parent) {
+ this.emitter.pipe(parent.emitter);
+ }
+ }
+
+ destroy() {
+ this.emitter.destroy();
+ this.controller.abort(new FrameworkError("Context destroyed."));
+ }
+
+ static enter(
+ self: RunInstance ,
+ fn: (context: RunContext ) => Promise,
+ signal?: AbortSignal,
+ ) {
+ const parent = RunContext.#storage.getStore();
+ const runContext = new RunContext(self, parent, signal);
+
+ return new Run(async () => {
+ const emitter = runContext.emitter.child({
+ namespace: ["run"],
+ creator: runContext,
+ context: { internal: true },
+ });
+
+ try {
+ await emitter.emit("start", null);
+ const result = await Promise.race([
+ RunContext.#storage.run(runContext, fn, runContext),
+ new Promise((_, reject) =>
+ runContext.signal.addEventListener("abort", () => reject(runContext.signal.reason)),
+ ),
+ ]);
+ await emitter.emit("success", result);
+ return result;
+ } catch (_e) {
+ const e = FrameworkError.ensure(_e);
+ await emitter.emit("error", e);
+ throw e;
+ } finally {
+ await emitter.emit("finish", null);
+ runContext.destroy();
+ }
+ }, runContext);
+ }
+
+ static {
+ this.register();
+ }
+
+ createSnapshot() {
+ return {
+ controller: this.controller,
+ runId: this.runId,
+ groupId: this.groupId,
+ parentId: this.parentId,
+ emitter: this.emitter,
+ context: this.context,
+ };
+ }
+
+ loadSnapshot(snapshot: ReturnType) {
+ Object.assign(this, snapshot);
+ }
+}
diff --git a/src/emitter/__snapshots__/emitter.test.ts.snap b/src/emitter/__snapshots__/emitter.test.ts.snap
new file mode 100644
index 00000000..9d138ddb
--- /dev/null
+++ b/src/emitter/__snapshots__/emitter.test.ts.snap
@@ -0,0 +1,36 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`Emitter > Complex workflow > Emits events in-order 1`] = `
+[
+ [
+ "start",
+ {
+ "a": 1,
+ },
+ ],
+ [
+ "success",
+ "a",
+ ],
+ [
+ "success",
+ "b",
+ ],
+ [
+ "start",
+ {
+ "a": 2,
+ },
+ ],
+ [
+ "start",
+ {
+ "a": 3,
+ },
+ ],
+ [
+ "success",
+ "c",
+ ],
+]
+`;
diff --git a/src/emitter/emitter.test.ts b/src/emitter/emitter.test.ts
new file mode 100644
index 00000000..75016b0c
--- /dev/null
+++ b/src/emitter/emitter.test.ts
@@ -0,0 +1,114 @@
+/**
+ * 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 { Emitter } from "@/emitter/emitter.js";
+import { describe, expect } from "vitest";
+import { EventMeta } from "@/emitter/types.js";
+
+describe("Emitter", () => {
+ it("Emits", async () => {
+ const emitter = new Emitter({
+ namespace: ["app"],
+ creator: this,
+ });
+ const fn = vitest.fn();
+ emitter.on("onStart", fn);
+ await emitter.emit("onStart", "1");
+ expect(fn).toBeCalledWith(
+ "1",
+ expect.objectContaining({
+ name: "onStart",
+ path: "app.onStart",
+ creator: this,
+ source: emitter,
+ createdAt: expect.any(Date),
+ context: expect.any(Object),
+ }),
+ );
+ });
+
+ describe("Complex workflow", () => {
+ const recorder: [keyof MyEvents, unknown][] = [];
+
+ interface MyEvents {
+ start: (data: { a: number }) => void;
+ success: (data: string) => void;
+ }
+
+ const emitter = new Emitter({
+ namespace: ["agent"],
+ });
+
+ it("Emits events in-order", async () => {
+ emitter.on("start", (data) => recorder.push(["start", data]));
+ emitter.on("success", (data) => recorder.push(["success", data]));
+ await emitter.emit("start", { a: 1 });
+ await emitter.emit("success", "a");
+ await emitter.emit("success", "b");
+ await emitter.emit("start", { a: 2 });
+ await emitter.emit("start", { a: 3 });
+ await emitter.emit("success", "c");
+
+ expect(recorder).toMatchSnapshot();
+ });
+
+ it("Resets", async () => {
+ recorder.length = 0;
+ emitter.reset();
+ await emitter.emit("success", "c");
+ expect(recorder).toHaveLength(0);
+ });
+ });
+
+ it("Accepts callbacks", async () => {
+ const emitter = new Emitter<{
+ start: (value: number) => void;
+ }>();
+ const recorder: string[] = [];
+ emitter.registerCallbacks({
+ start: (_, event) => recorder.push(event.path),
+ });
+ await emitter.emit("start", 10);
+ expect(recorder).toHaveLength(1);
+ });
+
+ it("Handles nesting", async () => {
+ const recorder: { source: string; meta: EventMeta }[] = [];
+
+ const root = Emitter.root;
+ root.match("*", (data, meta) => {
+ recorder.push({ source: "root.*", meta });
+ });
+ //root.on("agent.onStart", (_, meta) =>
+ // recorder.push({ name: meta.name, path: meta.path, source: "root" }),
+ //);
+ //root.on("agent.runner.onStart", (_, meta) => recorder.push({ source: "root.runner", meta }));
+
+ const agent = root.child({
+ namespace: ["agent"],
+ });
+ agent.on("onStart", (_, meta) => recorder.push({ source: "agent", meta }));
+ agent.on("*", (_, meta) => recorder.push({ source: "agent.*", meta }));
+ await agent.emit("onStart", null);
+
+ // agent.runner.onStart
+ const runner = agent.child({
+ namespace: ["runner"],
+ });
+ // agent.runner, agent.runner.onStart
+ await runner.emit("onStart", null);
+ });
+});
diff --git a/src/emitter/emitter.ts b/src/emitter/emitter.ts
new file mode 100644
index 00000000..1496ed74
--- /dev/null
+++ b/src/emitter/emitter.ts
@@ -0,0 +1,238 @@
+/**
+ * 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 {
+ Callback,
+ CleanupFn,
+ EmitterOptions,
+ EventMeta,
+ EventTrace,
+ Matcher,
+ StringKey,
+} from "@/emitter/types.js";
+export type { EventMeta, EventTrace };
+import { Cache } from "@/cache/decoratorCache.js";
+import { createFullPath, isPath, assertValidName, assertValidNamespace } from "@/emitter/utils.js";
+import { EmitterError } from "@/emitter/errors.js";
+import { createRandomHash } from "@/internals/helpers/hash.js";
+import { shallowCopy, toBoundedFunction } from "@/serializer/utils.js";
+import { RequiredAll } from "@/internals/types.js";
+import { Serializable } from "@/internals/serializable.js";
+import { pick } from "remeda";
+
+export interface EmitterInput {
+ groupId?: string;
+ namespace?: string[];
+ creator?: object;
+ context?: E & object;
+ trace?: EventTrace;
+}
+export interface EmitterChildInput {
+ groupId?: string;
+ namespace?: string[];
+ creator?: object;
+ context?: E & object;
+ trace?: EventTrace;
+}
+
+interface Listener {
+ readonly match: (event: EventMeta) => boolean;
+ readonly raw: Matcher;
+ readonly callback: Callback;
+ readonly options?: EmitterOptions;
+}
+
+export class Emitter>> extends Serializable {
+ protected listeners = new Set();
+ protected readonly groupId?: string;
+ public readonly namespace: string[];
+ public readonly creator?: object;
+ public readonly context: object;
+ protected readonly trace?: EventTrace;
+
+ constructor(input: EmitterInput = {}) {
+ super();
+ this.groupId = input?.groupId;
+ this.namespace = input?.namespace ?? [];
+ this.creator = input.creator ?? Object.prototype;
+ this.context = input?.context ?? {};
+ this.trace = input.trace;
+ assertValidNamespace(this.namespace);
+ }
+
+ static {
+ this.register();
+ }
+
+ @Cache()
+ static get root() {
+ return new Emitter>>({ creator: Object.create(null) });
+ }
+
+ child(input: EmitterChildInput = {}): Emitter {
+ const child = new Emitter({
+ trace: input.trace ?? this.trace,
+ groupId: input?.groupId ?? this.groupId,
+ context: { ...this.context, ...input?.context },
+ creator: input?.creator ?? this.creator,
+ namespace: input?.namespace
+ ? [...this.namespace, ...input.namespace]
+ : this.namespace.slice(),
+ });
+
+ child.pipe(this);
+ return child;
+ }
+
+ pipe(target: Emitter): CleanupFn {
+ return this.on(
+ // @ts-expect-error
+ "*.*",
+ toBoundedFunction(
+ // @ts-expect-error
+ (...args) => target.invoke(...args),
+ [{ value: target, name: "target" }],
+ ),
+ {
+ isBlocking: true,
+ once: false,
+ persistent: true,
+ },
+ );
+ }
+
+ destroy() {
+ this.listeners.clear();
+ }
+
+ reset() {
+ for (const listener of this.listeners) {
+ if (!listener.options?.persistent) {
+ this.listeners.delete(listener);
+ }
+ }
+ }
+
+ registerCallbacks>>(
+ callbacks: Partial<
+ Record extends Callback ? Callback : never>
+ >,
+ options?: Partial>,
+ ): CleanupFn {
+ const listeners: CleanupFn[] = [];
+ Object.entries(callbacks).forEach(([key, value]) => {
+ if (value) {
+ // @ts-expect-error
+ listeners.push(this.on(key, value, options?.[key]));
+ }
+ });
+ return () => listeners.forEach((cleanup) => cleanup());
+ }
+
+ on>>(
+ event: K,
+ callback: NonNullable extends Callback ? T[K] : never,
+ options?: EmitterOptions,
+ ): CleanupFn {
+ return this.match(event as Matcher, callback as Callback, options);
+ }
+
+ match(matcher: Matcher, callback: Callback, options?: EmitterOptions): CleanupFn {
+ const listener: Listener = {
+ options,
+ callback,
+ raw: matcher,
+ match: (() => {
+ if (matcher === "*") {
+ return (event) => event.path === createFullPath(this.namespace, event.name);
+ } else if (matcher === "*.*") {
+ return () => true;
+ } else if (matcher instanceof RegExp) {
+ return (event) => matcher.test(event.path);
+ } else if (typeof matcher === "function") {
+ return matcher;
+ } else if (typeof matcher === "string") {
+ return isPath(matcher)
+ ? (event) => event.path === matcher
+ : (event) =>
+ event.name === matcher && event.path === createFullPath(this.namespace, event.name);
+ } else {
+ throw new EmitterError("Invalid matcher provided!");
+ }
+ })(),
+ };
+ this.listeners.add(listener);
+
+ return () => this.listeners.delete(listener);
+ }
+
+ async emit>>(
+ name: K,
+ value: NonNullable extends Callback ? X : unknown,
+ ): Promise {
+ assertValidName(name);
+
+ const event = this.createEvent(name);
+ return await this.invoke(value, event);
+ }
+
+ protected async invoke(data: unknown, event: EventMeta) {
+ const executions: unknown[] = [];
+ for (const listener of this.listeners.values()) {
+ if (!listener.match(event)) {
+ continue;
+ }
+
+ if (listener.options?.once) {
+ this.listeners.delete(listener);
+ }
+
+ const run = async () => listener.callback(data, event);
+ executions.push(listener.options?.isBlocking ? await run() : run());
+ }
+ await Promise.all(executions);
+ }
+
+ protected createEvent(name: string): EventMeta {
+ return {
+ id: createRandomHash(),
+ groupId: this.groupId,
+ name,
+ path: createFullPath(this.namespace, name),
+ createdAt: new Date(),
+ source: this,
+ creator: this.creator!,
+ context: Object.assign({}, this.context, {}), // TODO: use createInStone
+ trace: shallowCopy(this.trace), // TODO
+ };
+ }
+
+ createSnapshot() {
+ return {
+ groupId: this.groupId,
+ namespace: shallowCopy(this.namespace),
+ creator: this.creator,
+ context: this.context,
+ trace: this.trace,
+ listeners: Array.from(this.listeners).map(pick(["raw", "options", "callback"])),
+ };
+ }
+
+ loadSnapshot({ listeners, ...snapshot }: ReturnType): void {
+ Object.assign(this, snapshot, { listeners: new Set() });
+ listeners.forEach(({ raw, callback, options }) => this.match(raw, callback, options));
+ }
+}
diff --git a/src/emitter/errors.ts b/src/emitter/errors.ts
new file mode 100644
index 00000000..57c2484e
--- /dev/null
+++ b/src/emitter/errors.ts
@@ -0,0 +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 { FrameworkError } from "@/errors.js";
+
+export class EmitterError extends FrameworkError {}
diff --git a/src/emitter/types.ts b/src/emitter/types.ts
new file mode 100644
index 00000000..473ddd9d
--- /dev/null
+++ b/src/emitter/types.ts
@@ -0,0 +1,56 @@
+/**
+ * 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 { AnyVoid } from "@/internals/types.js";
+import { Emitter } from "@/emitter/emitter.js";
+
+export type Matcher = "*" | "*.*" | RegExp | ((event: EventMeta) => boolean);
+//export type Callback