release v0.2.9

Co-authored-by: Yi Ding <yi.s.ding@gmail.com>

.changeset/fast-mayflies-grab.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-gpt-4-turbo GA

.changeset/late-wombats-travel.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Mistral 0.1.3

.eslintrc.cjs

@@ -1,12 +1,22 @@
 module.exports = {
+  root: true,
   extends: [
-    "next",
     "turbo",
     "prettier",
     "plugin:@typescript-eslint/recommended-type-checked-only",
   ],
+  parserOptions: {
+    project: true,
+    __tsconfigRootDir: __dirname,
+  },
+  settings: {
+    react: {
+      version: "999.999.999",
+    },
+  },
   rules: {
-    "@next/next/no-html-link-for-pages": "off",
+    "max-params": ["error", 4],
+    "prefer-const": "error",
     "@typescript-eslint/no-floating-promises": [
       "error",
       {
@@ -54,22 +64,13 @@ module.exports = {
     "@typescript-eslint/triple-slash-reference": "off",
     "@typescript-eslint/unbound-method": "off",
   },
-  // NOTE I think because we've temporarily removed all of the NextJS stuff
-  // from the turborepo not having next in the devDeps causes an error on only
-  // clean clones of the repo
-  // Not sure if this is a missing dependency in the package.json or just my not
-  // understanding how turborepo is supposed to work.
-  // Anyways, planning to add back a Next.JS example soon
-  parserOptions: {
-    babelOptions: {
-      presets: [require.resolve("next/babel")],
+  overrides: [
+    {
+      files: ["examples/**/*.ts"],
+      rules: {
+        "turbo/no-undeclared-env-vars": "off",
+      },
     },
-    project: true,
-    __tsconfigRootDir: __dirname,
-  },
-  settings: {
-    react: {
-      version: "999.999.999",
-    },
-  },
+  ],
+  ignorePatterns: ["dist/", "lib/"],
 };

.eslintrc.js

@@ -1,23 +0,0 @@
-module.exports = {
-  root: true,
-  // This tells ESLint to load the config from the package `eslint-config-custom`
-  extends: ["custom"],
-  settings: {
-    next: {
-      rootDir: ["apps/*/"],
-    },
-  },
-  rules: {
-    "max-params": ["error", 4],
-    "prefer-const": "error",
-  },
-  overrides: [
-    {
-      files: ["examples/**/*.ts"],
-      rules: {
-        "turbo/no-undeclared-env-vars": "off",
-      },
-    },
-  ],
-  ignorePatterns: ["dist/", "lib/"],
-};

.github/workflows/lint_on_push_or_pull.yml

@@ -14,8 +14,6 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: pnpm/action-setup@v2
-        with:
-          version: latest
       - name: Setup Node.js
         uses: actions/setup-node@v4
         with:

.github/workflows/test.yml

@@ -1,6 +1,12 @@
 name: Run Tests
 
-on: [push, pull_request]
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
@@ -8,6 +14,11 @@ concurrency:
 
 jobs:
   e2e:
+    strategy:
+      fail-fast: false
+      matrix:
+        node-version: [18.x, 20.x, 21.x]
+    name: E2E on Node.js ${{ matrix.node-version }}
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -57,6 +68,9 @@ jobs:
         run: pnpm install
       - name: Build
         run: pnpm run build --filter llamaindex
+      - name: Use Build For Examples
+        run: pnpm link ../packages/core/
+        working-directory: ./examples
       - name: Run Type Check
         run: pnpm run type-check
       - name: Run Circular Dependency Check

.vscode/launch.json

@@ -10,8 +10,9 @@
       "name": "Debug Example",
       "skipFiles": ["<node_internals>/**"],
       "runtimeExecutable": "pnpm",
+      "console": "integratedTerminal",
       "cwd": "${workspaceFolder}/examples",
-      "runtimeArgs": ["ts-node", "${fileBasename}"]
+      "runtimeArgs": ["npx", "tsx", "${file}"]
     }
   ]
 }

README.md

@@ -154,7 +154,7 @@ If you need any of those classes, you have to import them instead directly. Here
 import { PineconeVectorStore } from "@llamaindex/edge/storage/vectorStore/PineconeVectorStore";
 ```
 
-As the `PDFReader` is not with the Edge runtime, here's how to use the `SimpleDirectoryReader` with the `LlamaParseReader` to load PDFs:
+As the `PDFReader` is not working with the Edge runtime, here's how to use the `SimpleDirectoryReader` with the `LlamaParseReader` to load PDFs:
 
 ```typescript
 import { SimpleDirectoryReader } from "@llamaindex/edge/readers/SimpleDirectoryReader";

apps/docs/docs/examples/agent.mdx

@@ -67,7 +67,6 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [sumFunctionTool, divideFunctionTool],
-    verbose: true,
   });
 
   // Chat with the agent

apps/docs/docs/modules/agent/multi_document_agent.mdx

@@ -221,7 +221,6 @@ for (const title of wikiTitles) {
   const agent = new OpenAIAgent({
     tools: queryEngineTools,
     llm,
-    verbose: true,
   });
 
   documentAgents[title] = agent;
@@ -282,7 +281,6 @@ const objectIndex = await ObjectIndex.fromObjects(
 const topAgent = new OpenAIAgent({
   toolRetriever: await objectIndex.asRetriever({}),
   llm,
-  verbose: true,
   prefixMessages: [
     {
       content:

apps/docs/docs/modules/agent/openai.mdx

@@ -88,7 +88,6 @@ Now we can create an OpenAIAgent with the function tools.
 ```ts
 const agent = new OpenAIAgent({
   tools: [sumFunctionTool, divideFunctionTool],
-  verbose: true,
 });
 ```
 
@@ -169,7 +168,6 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [sumFunctionTool, divideFunctionTool],
-    verbose: true,
   });
 
   // Chat with the agent

apps/docs/docs/modules/agent/query_engine_tool.mdx

@@ -64,7 +64,6 @@ const queryEngineTool = new QueryEngineTool({
 
 const agent = new OpenAIAgent({
   tools: [queryEngineTool],
-  verbose: true,
 });
 ```
 
@@ -114,7 +113,6 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [queryEngineTool],
-    verbose: true,
   });
 
   // Chat with the agent

apps/docs/docs/modules/agent/react_agent.mdx

@@ -90,7 +90,6 @@ Now we can create an OpenAIAgent with the function tools.
 ```ts
 const agent = new ReActAgent({
   tools: [sumFunctionTool, divideFunctionTool],
-  verbose: true,
 });
 ```
 
@@ -185,7 +184,6 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [sumFunctionTool, divideFunctionTool],
-    verbose: true,
   });
 
   // Chat with the agent

apps/docs/docs/modules/llms/available_llms/llama2.md

@@ -3,7 +3,7 @@
 ## Usage
 
 ```ts
-import { Ollama, Settings } from "llamaindex";
+import { Ollama, Settings, DeuceChatStrategy } from "llamaindex";
 
 Settings.llm = new LlamaDeuce({ chatStrategy: DeuceChatStrategy.META });
 ```
@@ -11,7 +11,12 @@ Settings.llm = new LlamaDeuce({ chatStrategy: DeuceChatStrategy.META });
 ## Usage with Replication
 
 ```ts
-import { Ollama, ReplicateSession, Settings } from "llamaindex";
+import {
+  Ollama,
+  ReplicateSession,
+  Settings,
+  DeuceChatStrategy,
+} from "llamaindex";
 
 const replicateSession = new ReplicateSession({
   replicateKey,
@@ -48,7 +53,13 @@ const results = await queryEngine.query({
 ## Full Example
 
 ```ts
-import { LlamaDeuce, Document, VectorStoreIndex, Settings } from "llamaindex";
+import {
+  LlamaDeuce,
+  Document,
+  VectorStoreIndex,
+  Settings,
+  DeuceChatStrategy,
+} from "llamaindex";
 
 // Use the LlamaDeuce LLM
 Settings.llm = new LlamaDeuce({ chatStrategy: DeuceChatStrategy.META });

apps/docs/docs/observability/index.md

@@ -14,6 +14,9 @@ Configure a variable once, and you'll be able to do things like the following:
 
 Each provider has similarities and differences. Take a look below for the full set of guides for each one!
 
+- [OpenLLMetry](#openllmetry)
+- [Langtrace](#langtrace)
+
 ## OpenLLMetry
 
 [OpenLLMetry](https://github.com/traceloop/openllmetry-js) is an open-source project based on OpenTelemetry for tracing and monitoring
@@ -33,3 +36,29 @@ traceloop.initialize({
   disableBatch: true,
 });
 ```
+
+## Langtrace
+
+Enhance your observability with Langtrace, a robust open-source tool supports OpenTelemetry and is designed to trace, evaluate, and manage LLM applications seamlessly. Langtrace integrates directly with LlamaIndex, offering detailed, real-time insights into performance metrics such as accuracy, evaluations, and latency.
+
+#### Install
+
+- Self-host or sign-up and generate an API key using [Langtrace](https://www.langtrace.ai) Cloud
+
+```bash
+npm install @langtrase/typescript-sdk
+```
+
+#### Initialize
+
+```js
+import * as Langtrace from "@langtrase/typescript-sdk";
+Langtrace.init({ api_key: "<YOUR_API_KEY>" });
+```
+
+Features:
+
+- OpenTelemetry compliant, ensuring broad compatibility with observability platforms.
+- Provides comprehensive logs and detailed traces of all components.
+- Real-time monitoring of accuracy, evaluations, usage, costs, and latency.
+- For more configuration options and details, visit [Langtrace Docs](https://docs.langtrace.ai/introduction).

apps/docs/docs/recipes/cost-analysis.mdx

@@ -6,7 +6,7 @@ This page shows how to track LLM cost using APIs.
 
 The callback manager is a class that manages the callback functions.
 
-You can register `llm-start`, and `llm-end` callbacks to the callback manager for tracking the cost.
+You can register `llm-start`, `llm-end`, and `llm-stream` callbacks to the callback manager for tracking the cost.
 
 import CodeBlock from "@theme/CodeBlock";
 import CodeSource from "!raw-loader!../../../../examples/recipes/cost-analysis";

apps/docs/package.json

@@ -32,7 +32,7 @@
     "@docusaurus/theme-classic": "^3.2.1",
     "@docusaurus/types": "^3.2.1",
     "@tsconfig/docusaurus": "^2.0.3",
-    "@types/node": "^18.19.31",
+    "@types/node": "^20.12.7",
     "docusaurus-plugin-typedoc": "^0.22.0",
     "typedoc": "^0.25.13",
     "typedoc-plugin-markdown": "^3.17.1",

examples/agent/openai.ts

@@ -1,13 +1,13 @@
 import { FunctionTool, OpenAIAgent } from "llamaindex";
 
 // Define a function to sum two numbers
-function sumNumbers({ a, b }: { a: number; b: number }): number {
-  return a + b;
+function sumNumbers({ a, b }: { a: number; b: number }) {
+  return `${a + b}`;
 }
 
 // Define a function to divide two numbers
-function divideNumbers({ a, b }: { a: number; b: number }): number {
-  return a / b;
+function divideNumbers({ a, b }: { a: number; b: number }) {
+  return `${a / b}`;
 }
 
 // Define the parameters of the sum function as a JSON schema
@@ -24,7 +24,7 @@ const sumJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 const divideJSON = {
   type: "object",
@@ -39,7 +39,7 @@ const divideJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 async function main() {
   // Create a function tool from the sum function

examples/agent/react_agent.ts

@@ -1,13 +1,13 @@
 import { Anthropic, FunctionTool, ReActAgent } from "llamaindex";
 
 // Define a function to sum two numbers
-function sumNumbers({ a, b }: { a: number; b: number }): number {
-  return a + b;
+function sumNumbers({ a, b }: { a: number; b: number }) {
+  return `${a + b}`;
 }
 
 // Define a function to divide two numbers
-function divideNumbers({ a, b }: { a: number; b: number }): number {
-  return a / b;
+function divideNumbers({ a, b }: { a: number; b: number }) {
+  return `${a / b}`;
 }
 
 // Define the parameters of the sum function as a JSON schema
@@ -24,7 +24,7 @@ const sumJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 const divideJSON = {
   type: "object",
@@ -39,7 +39,7 @@ const divideJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 async function main() {
   // Create a function tool from the sum function

examples/agent/step_wise_openai.ts

@@ -1,13 +1,13 @@
 import { FunctionTool, OpenAIAgent } from "llamaindex";
 
 // Define a function to sum two numbers
-function sumNumbers({ a, b }: { a: number; b: number }): number {
-  return a + b;
+function sumNumbers({ a, b }: { a: number; b: number }) {
+  return `${a + b}`;
 }
 
 // Define a function to divide two numbers
-function divideNumbers({ a, b }: { a: number; b: number }): number {
-  return a / b;
+function divideNumbers({ a, b }: { a: number; b: number }) {
+  return `${a / b}`;
 }
 
 // Define the parameters of the sum function as a JSON schema
@@ -24,22 +24,22 @@ const sumJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 const divideJSON = {
   type: "object",
   properties: {
     a: {
       type: "number",
-      description: "The dividend a to divide",
+      description: "The dividend",
     },
     b: {
       type: "number",
-      description: "The divisor b to divide by",
+      description: "The divisor",
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 async function main() {
   // Create a function tool from the sum function

examples/agent/step_wise_react.ts

@@ -1,13 +1,13 @@
 import { FunctionTool, ReActAgent } from "llamaindex";
 
 // Define a function to sum two numbers
-function sumNumbers({ a, b }: { a: number; b: number }): number {
-  return a + b;
+function sumNumbers({ a, b }: { a: number; b: number }) {
+  return `${a + b}`;
 }
 
 // Define a function to divide two numbers
-function divideNumbers({ a, b }: { a: number; b: number }): number {
-  return a / b;
+function divideNumbers({ a, b }: { a: number; b: number }) {
+  return `${a / b}`;
 }
 
 // Define the parameters of the sum function as a JSON schema
@@ -24,7 +24,7 @@ const sumJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 const divideJSON = {
   type: "object",
@@ -39,7 +39,7 @@ const divideJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 async function main() {
   // Create a function tool from the sum function

examples/agent/stream_openai_agent.ts

@@ -1,13 +1,13 @@
 import { FunctionTool, OpenAIAgent } from "llamaindex";
 
 // Define a function to sum two numbers
-function sumNumbers({ a, b }: { a: number; b: number }): number {
-  return a + b;
+function sumNumbers({ a, b }: { a: number; b: number }) {
+  return `${a + b}`;
 }
 
 // Define a function to divide two numbers
-function divideNumbers({ a, b }: { a: number; b: number }): number {
-  return a / b;
+function divideNumbers({ a, b }: { a: number; b: number }) {
+  return `${a / b}`;
 }
 
 // Define the parameters of the sum function as a JSON schema
@@ -24,7 +24,7 @@ const sumJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 const divideJSON = {
   type: "object",
@@ -39,18 +39,18 @@ const divideJSON = {
     },
   },
   required: ["a", "b"],
-};
+} as const;
 
 async function main() {
   // Create a function tool from the sum function
-  const functionTool = new FunctionTool(sumNumbers, {
+  const functionTool = FunctionTool.from(sumNumbers, {
     name: "sumNumbers",
     description: "Use this function to sum two numbers",
     parameters: sumJSON,
   });
 
   // Create a function tool from the divide function
-  const functionTool2 = new FunctionTool(divideNumbers, {
+  const functionTool2 = FunctionTool.from(divideNumbers, {
     name: "divideNumbers",
     description: "Use this function to divide two numbers",
     parameters: divideJSON,

examples/anthropic/agent.ts

@@ -0,0 +1,43 @@
+import { FunctionTool, Settings, WikipediaTool } from "llamaindex";
+import { AnthropicAgent } from "llamaindex/agent/anthropic";
+
+Settings.callbackManager.on("llm-tool-call", (event) => {
+  console.log("llm-tool-call", event.detail.payload.toolCall);
+});
+
+const agent = new AnthropicAgent({
+  tools: [
+    FunctionTool.from<{ location: string }>(
+      (query) => {
+        return `The weather in ${query.location} is sunny`;
+      },
+      {
+        name: "weather",
+        description: "Get the weather",
+        parameters: {
+          type: "object",
+          properties: {
+            location: {
+              type: "string",
+              description: "The location to get the weather for",
+            },
+          },
+          required: ["location"],
+        },
+      },
+    ),
+    new WikipediaTool(),
+  ],
+});
+
+async function main() {
+  // https://docs.anthropic.com/claude/docs/tool-use#tool-use-best-practices-and-limitations
+  const { response } = await agent.chat({
+    message:
+      "What is the weather in New York? What's the history of New York from Wikipedia in 3 sentences?",
+  });
+
+  console.log(response);
+}
+
+void main();

examples/multimodal/load.ts

@@ -4,6 +4,7 @@ import {
   VectorStoreIndex,
   storageContextFromDefaults,
 } from "llamaindex";
+import { DocStoreStrategy } from "llamaindex/ingestion/strategies/index";
 
 import * as path from "path";
 
@@ -31,6 +32,7 @@ async function generateDatasource() {
     });
     await VectorStoreIndex.fromDocuments(documents, {
       storageContext,
+      docStoreStrategy: DocStoreStrategy.NONE,
     });
   });
   console.log(`Storage successfully generated in ${ms / 1000}s.`);

examples/package.json

@@ -5,21 +5,22 @@
   "dependencies": {
     "@aws-crypto/sha256-js": "^5.2.0",
     "@datastax/astra-db-ts": "^0.1.4",
-    "@notionhq/client": "^2.2.14",
+    "@notionhq/client": "^2.2.15",
     "@pinecone-database/pinecone": "^1.1.3",
     "@zilliz/milvus2-sdk-node": "^2.3.5",
     "chromadb": "^1.8.1",
     "commander": "^11.1.0",
     "dotenv": "^16.4.5",
     "js-tiktoken": "^1.0.10",
-    "llamaindex": "workspace:latest",
+    "llamaindex": "latest",
     "mongodb": "^6.5.0",
     "pathe": "^1.1.2"
   },
   "devDependencies": {
-    "@types/node": "^18.19.31",
+    "@types/node": "^20.12.7",
     "ts-node": "^10.9.2",
-    "typescript": "^5.4.4"
+    "tsx": "^4.7.2",
+    "typescript": "^5.4.5"
   },
   "scripts": {
     "lint": "eslint ."

examples/readers/package.json

@@ -15,7 +15,7 @@
     "llamaindex": "latest"
   },
   "devDependencies": {
-    "@types/node": "^20.11.14",
+    "@types/node": "^20.12.7",
     "ts-node": "^10.9.2",
     "typescript": "^5.4.3"
   }

examples/recipes/cost-analysis.ts

@@ -6,7 +6,8 @@ import { extractText } from "llamaindex/llm/utils";
 const encoding = encodingForModel("gpt-4-0125-preview");
 
 const llm = new OpenAI({
-  model: "gpt-4-0125-preview",
+  // currently is "gpt-4-turbo-2024-04-09"
+  model: "gpt-4-turbo",
 });
 
 let tokenCount = 0;
@@ -19,18 +20,25 @@ Settings.callbackManager.on("llm-start", (event) => {
   console.log("Token count:", tokenCount);
   // https://openai.com/pricing
   // $10.00 / 1M tokens
-  console.log(`Price: $${(tokenCount / 1_000_000) * 10}`);
-});
-Settings.callbackManager.on("llm-end", (event) => {
-  const { response } = event.detail.payload;
-  tokenCount += encoding.encode(extractText(response.message.content)).length;
-  console.log("Token count:", tokenCount);
-  // https://openai.com/pricing
-  // $30.00 / 1M tokens
-  console.log(`Price: $${(tokenCount / 1_000_000) * 30}`);
+  console.log(`Total Price: $${(tokenCount / 1_000_000) * 10}`);
 });
 
-const question = "Hello, how are you?";
+Settings.callbackManager.on("llm-stream", (event) => {
+  const { chunk } = event.detail.payload;
+  const { delta } = chunk;
+  tokenCount += encoding.encode(extractText(delta)).length;
+  if (tokenCount > 20) {
+    // This is just an example, you can set your own limit or handle it differently
+    throw new Error("Token limit exceeded!");
+  }
+});
+Settings.callbackManager.on("llm-end", () => {
+  // https://openai.com/pricing
+  // $30.00 / 1M tokens
+  console.log(`Total Price: $${(tokenCount / 1_000_000) * 30}`);
+});
+
+const question = "Hello, how are you? Please response about 50 tokens.";
 console.log("Question:", question);
 void llm
   .chat({

examples/toolsStream.ts

@@ -1,8 +1,7 @@
-import { ChatResponseChunk, OpenAI } from "llamaindex";
+import { OpenAI } from "llamaindex";
 
 async function main() {
   const llm = new OpenAI({ model: "gpt-4-turbo" });
-
   const args: Parameters<typeof llm.chat>[0] = {
     additionalChatOptions: {
       tool_choice: "auto",
@@ -34,11 +33,13 @@ async function main() {
   };
 
   const stream = await llm.chat({ ...args, stream: true });
-  let chunk: ChatResponseChunk | null = null;
-  for await (chunk of stream) {
+  for await (const chunk of stream) {
     process.stdout.write(chunk.delta);
+    if (chunk.options && "toolCall" in chunk.options) {
+      console.log("Tool call:");
+      console.log(chunk.options.toolCall);
+    }
   }
-  console.log(chunk?.additionalKwargs?.toolCalls[0]);
 }
 
 (async function () {

examples/vectorIndex.ts

@@ -1,6 +1,11 @@
 import fs from "node:fs/promises";
 
-import { Document, VectorStoreIndex } from "llamaindex";
+import {
+  Document,
+  MetadataMode,
+  NodeWithScore,
+  VectorStoreIndex,
+} from "llamaindex";
 
 async function main() {
   // Load essay from abramov.txt in Node
@@ -16,12 +21,20 @@ async function main() {
 
   // Query the index
   const queryEngine = index.asQueryEngine();
-  const response = await queryEngine.query({
+  const { response, sourceNodes } = await queryEngine.query({
     query: "What did the author do in college?",
   });
 
-  // Output response
-  console.log(response.toString());
+  // Output response with sources
+  console.log(response);
+
+  if (sourceNodes) {
+    sourceNodes.forEach((source: NodeWithScore, index: number) => {
+      console.log(
+        `\n${index}: Score: ${source.score} - ${source.node.getContent(MetadataMode.NONE).substring(0, 50)}...\n`,
+      );
+    });
+  }
 }
 
 main().catch(console.error);

package.json

@@ -15,25 +15,30 @@
     "release": "pnpm run check-minor-version && pnpm run build:release && changeset publish",
     "release-snapshot": "pnpm run check-minor-version && pnpm run build:release && changeset publish --tag snapshot",
     "check-minor-version": "node ./scripts/check-minor-version",
-    "new-version": "pnpm run build:release && changeset version && pnpm run check-minor-version",
+    "new-version": "changeset version && pnpm run check-minor-version && pnpm run build:release",
     "new-snapshot": "pnpm run build:release && changeset version --snapshot"
   },
   "devDependencies": {
     "@changesets/cli": "^2.27.1",
+    "@typescript-eslint/eslint-plugin": "^7.7.0",
     "eslint": "^8.57.0",
-    "eslint-config-custom": "workspace:*",
+    "eslint-config-next": "^13.5.6",
+    "eslint-config-prettier": "^8.10.0",
+    "eslint-config-turbo": "^1.13.2",
+    "eslint-plugin-react": "7.28.0",
     "husky": "^9.0.11",
     "lint-staged": "^15.2.2",
     "prettier": "^3.2.5",
     "prettier-plugin-organize-imports": "^3.2.4",
     "turbo": "^1.13.2",
-    "typescript": "^5.4.4"
+    "typescript": "^5.4.5"
   },
-  "packageManager": "pnpm@8.15.6+sha256.01c01eeb990e379b31ef19c03e9d06a14afa5250b82e81303f88721c99ff2e6f",
+  "packageManager": "pnpm@9.0.1+sha256.46d50ee2afecb42b185ebbd662dc7bdd52ef5be56bf035bb615cab81a75345df",
   "pnpm": {
     "overrides": {
       "trim": "1.0.1",
-      "@babel/traverse": "7.23.2"
+      "@babel/traverse": "7.23.2",
+      "protobufjs": "7.2.6"
     }
   },
   "lint-staged": {

packages/core/CHANGELOG.md

@@ -1,5 +1,49 @@
 # llamaindex
 
+## 0.2.9
+
+### Patch Changes
+
+- 76c3fd6: Add score to source nodes response
+- 208282d: feat: init anthropic agent
+
+  remove the `tool` | `function` type in `MessageType`. Replace with `assistant` instead.
+  This is because these two types are only available for `OpenAI`.
+  Since `OpenAI` deprecates the function type, we support the Claude 3 tool call.
+
+## 0.2.8
+
+### Patch Changes
+
+- Add ToolsFactory to generate agent tools
+
+## 0.2.7
+
+### Patch Changes
+
+- 96f8f40: fix: agent stream
+- Updated dependencies
+  - @llamaindex/env@0.0.7
+
+## 0.2.6
+
+### Patch Changes
+
+- a3b4409: Fix agent streaming with new OpenAI models
+
+## 0.2.5
+
+### Patch Changes
+
+- 7d56cdf: Allow OpenAIAgent to be called without tools
+
+## 0.2.4
+
+### Patch Changes
+
+- 3bc77f7: gpt-4-turbo GA
+- 8d2b21e: Mistral 0.1.3
+
 ## 0.2.3
 
 ### Patch Changes

packages/core/e2e/fixtures/embeddings/OpenAIEmbedding.ts

@@ -0,0 +1,34 @@
+import { BaseNode, SimilarityType, type BaseEmbedding } from "llamaindex";
+
+export class OpenAIEmbedding implements BaseEmbedding {
+  embedBatchSize = 512;
+
+  async getQueryEmbedding(query: string) {
+    return [0];
+  }
+
+  async getTextEmbedding(text: string) {
+    return [0];
+  }
+
+  async getTextEmbeddings(texts: string[]) {
+    return [[0]];
+  }
+
+  async getTextEmbeddingsBatch(texts: string[]) {
+    return [[0]];
+  }
+
+  similarity(
+    embedding1: number[],
+    embedding2: number[],
+    mode?: SimilarityType,
+  ) {
+    return 1;
+  }
+
+  async transform(nodes: BaseNode[], _options?: any): Promise<BaseNode[]> {
+    nodes.forEach((node) => (node.embedding = [0]));
+    return nodes;
+  }
+}

packages/core/e2e/fixtures/llm/anthropic.ts

@@ -0,0 +1,3 @@
+import { OpenAI } from "./open_ai.js";
+
+export class Anthropic extends OpenAI {}

packages/core/e2e/fixtures/llm/open_ai.ts

@@ -1,4 +1,3 @@
-import { faker } from "@faker-js/faker";
 import type {
   ChatResponse,
   ChatResponseChunk,
@@ -9,6 +8,9 @@ import type {
   LLMCompletionParamsNonStreaming,
   LLMCompletionParamsStreaming,
 } from "llamaindex/llm/types";
+import { extractText } from "llamaindex/llm/utils";
+import { deepStrictEqual, strictEqual } from "node:assert";
+import { llmCompleteMockStorage } from "../../node/utils.js";
 
 export function getOpenAISession() {
   return {};
@@ -19,6 +21,7 @@ export function isFunctionCallingModel() {
 }
 
 export class OpenAI implements LLM {
+  supportToolCall = true;
   get metadata() {
     return {
       model: "mock-model",
@@ -40,21 +43,39 @@ export class OpenAI implements LLM {
       | LLMChatParamsStreaming<Record<string, unknown>>
       | LLMChatParamsNonStreaming<Record<string, unknown>>,
   ): unknown {
-    if (params.stream) {
-      return {
-        [Symbol.asyncIterator]: async function* () {
-          yield {
-            delta: faker.word.words(),
-          } satisfies ChatResponseChunk;
-        },
-      };
+    if (llmCompleteMockStorage.llmEventStart.length > 0) {
+      const chatMessage =
+        llmCompleteMockStorage.llmEventStart.shift()!["messages"];
+      strictEqual(chatMessage.length, params.messages.length);
+      for (let i = 0; i < chatMessage.length; i++) {
+        strictEqual(chatMessage[i].role, params.messages[i].role);
+        deepStrictEqual(chatMessage[i].content, params.messages[i].content);
+      }
+
+      if (llmCompleteMockStorage.llmEventEnd.length > 0) {
+        const { id, response } = llmCompleteMockStorage.llmEventEnd.shift()!;
+        if (params.stream) {
+          return {
+            [Symbol.asyncIterator]: async function* () {
+              while (true) {
+                const idx = llmCompleteMockStorage.llmEventStream.findIndex(
+                  (e) => e.id === id,
+                );
+                if (idx === -1) {
+                  break;
+                }
+                const chunk = llmCompleteMockStorage.llmEventStream[idx].chunk;
+                llmCompleteMockStorage.llmEventStream.splice(idx, 1);
+                yield chunk;
+              }
+            },
+          };
+        } else {
+          return response;
+        }
+      }
     }
-    return {
-      message: {
-        content: faker.lorem.paragraph(),
-        role: "assistant",
-      },
-    } satisfies ChatResponse;
+    throw new Error("Method not implemented.");
   }
   complete(
     params: LLMCompletionParamsStreaming,
@@ -62,7 +83,23 @@ export class OpenAI implements LLM {
   complete(
     params: LLMCompletionParamsNonStreaming,
   ): Promise<CompletionResponse>;
-  async complete(params: unknown): Promise<unknown> {
+  async complete(
+    params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
+  ): Promise<AsyncIterable<CompletionResponse> | CompletionResponse> {
+    if (llmCompleteMockStorage.llmEventStart.length > 0) {
+      const chatMessage =
+        llmCompleteMockStorage.llmEventStart.shift()!["messages"];
+      strictEqual(chatMessage.length, 1);
+      strictEqual(chatMessage[0].role, "user");
+      strictEqual(chatMessage[0].content, params.prompt);
+    }
+    if (llmCompleteMockStorage.llmEventEnd.length > 0) {
+      const response = llmCompleteMockStorage.llmEventEnd.shift()!["response"];
+      return {
+        raw: response,
+        text: extractText(response.message.content),
+      } satisfies CompletionResponse;
+    }
     throw new Error("Method not implemented.");
   }
 }

packages/core/e2e/node/basic.e2e.ts

@@ -1,139 +0,0 @@
-/* eslint-disable @typescript-eslint/no-floating-promises */
-import { consola } from "consola";
-import {
-  OpenAI,
-  OpenAIAgent,
-  Settings,
-  type LLM,
-  type LLMEndEvent,
-  type LLMStartEvent,
-} from "llamaindex";
-import { ok } from "node:assert";
-import type { WriteStream } from "node:fs";
-import { createWriteStream } from "node:fs";
-import { mkdir } from "node:fs/promises";
-import { join } from "node:path";
-import { after, before, beforeEach, describe, test } from "node:test";
-import { inspect } from "node:util";
-
-let llm: LLM;
-let fsStream: WriteStream;
-before(async () => {
-  const logUrl = new URL(
-    join(
-      "..",
-      "logs",
-      `basic.e2e.${new Date().toISOString().replace(/:/g, "-").replace(/\./g, "-")}.log`,
-    ),
-    import.meta.url,
-  );
-  await mkdir(new URL(".", logUrl), { recursive: true });
-  fsStream = createWriteStream(logUrl, {
-    encoding: "utf-8",
-  });
-});
-
-after(() => {
-  fsStream.end();
-});
-
-beforeEach((s) => {
-  fsStream.write("start: " + s.name + "\n");
-});
-
-const llmEventStartHandler = (event: LLMStartEvent) => {
-  const { payload } = event.detail;
-  fsStream.write(
-    "llmEventStart: " +
-      inspect(payload, {
-        depth: Infinity,
-      }) +
-      "\n",
-  );
-};
-
-const llmEventEndHandler = (event: LLMEndEvent) => {
-  const { payload } = event.detail;
-  fsStream.write(
-    "llmEventEnd: " +
-      inspect(payload, {
-        depth: Infinity,
-      }) +
-      "\n",
-  );
-};
-
-before(() => {
-  Settings.llm = new OpenAI({
-    model: "gpt-3.5-turbo",
-  });
-  llm = Settings.llm;
-  Settings.callbackManager.on("llm-start", llmEventStartHandler);
-  Settings.callbackManager.on("llm-end", llmEventEndHandler);
-});
-
-after(() => {
-  Settings.callbackManager.off("llm-start", llmEventStartHandler);
-  Settings.callbackManager.off("llm-end", llmEventEndHandler);
-});
-
-describe("llm", () => {
-  test("llm.chat", async () => {
-    const response = await llm.chat({
-      messages: [
-        {
-          content: "Hello",
-          role: "user",
-        },
-      ],
-    });
-    consola.debug("response:", response);
-    ok(typeof response.message.content === "string");
-  });
-
-  test("stream llm.chat", async () => {
-    const iter = await llm.chat({
-      stream: true,
-      messages: [
-        {
-          content: "hello",
-          role: "user",
-        },
-      ],
-    });
-    for await (const chunk of iter) {
-      consola.debug("chunk:", chunk);
-      ok(typeof chunk.delta === "string");
-    }
-  });
-});
-
-describe("agent", () => {
-  test("agent.chat", async () => {
-    const agent = new OpenAIAgent({
-      tools: [
-        {
-          call: async () => {
-            return "35 degrees and sunny in San Francisco";
-          },
-          metadata: {
-            name: "Weather",
-            description: "Get the weather",
-            parameters: {
-              type: "object",
-              properties: {
-                location: { type: "string" },
-              },
-              required: ["location"],
-            },
-          },
-        },
-      ],
-    });
-    const result = await agent.chat({
-      message: "What is the weather in San Francisco?",
-    });
-    consola.debug("response:", result.response);
-    ok(typeof result.response === "string");
-  });
-});

packages/core/e2e/node/claude.e2e.ts

@@ -0,0 +1,129 @@
+import { consola } from "consola";
+import { Anthropic, FunctionTool, Settings, type LLM } from "llamaindex";
+import { AnthropicAgent } from "llamaindex/agent/anthropic";
+import { ok } from "node:assert";
+import { beforeEach, test } from "node:test";
+import { sumNumbersTool } from "./fixtures/tools.js";
+import { mockLLMEvent } from "./utils.js";
+
+let llm: LLM;
+beforeEach(async () => {
+  Settings.llm = new Anthropic({
+    model: "claude-3-opus",
+  });
+  llm = Settings.llm;
+});
+
+await test("anthropic llm", async (t) => {
+  await mockLLMEvent(t, "llm-anthropic");
+  await t.test("llm.chat", async () => {
+    const response = await llm.chat({
+      messages: [
+        {
+          content: "Hello",
+          role: "user",
+          options: {},
+        },
+      ],
+    });
+    consola.debug("response:", response);
+    ok(typeof response.message.content === "string");
+  });
+
+  await t.test("stream llm.chat", async () => {
+    const iter = await llm.chat({
+      stream: true,
+      messages: [
+        {
+          content: "hello",
+          role: "user",
+        },
+      ],
+    });
+    for await (const chunk of iter) {
+      consola.debug("chunk:", chunk);
+      ok(typeof chunk.delta === "string");
+    }
+  });
+});
+
+await test("anthropic agent", async (t) => {
+  await mockLLMEvent(t, "anthropic-agent");
+  await t.test("chat", async () => {
+    const agent = new AnthropicAgent({
+      tools: [
+        {
+          call: async () => {
+            return "35 degrees and sunny in San Francisco";
+          },
+          metadata: {
+            name: "Weather",
+            description: "Get the weather",
+            parameters: {
+              type: "object",
+              properties: {
+                location: { type: "string" },
+              },
+              required: ["location"],
+            },
+          },
+        },
+      ],
+    });
+    const result = await agent.chat({
+      message: "What is the weather in San Francisco?",
+    });
+    consola.debug("response:", result.response);
+    ok(typeof result.response === "string");
+    ok(result.response.includes("35"));
+  });
+
+  await t.test("async function", async () => {
+    const uniqueId = "123456789";
+    const showUniqueId = FunctionTool.from<{
+      firstName: string;
+      lastName: string;
+    }>(
+      async ({ firstName, lastName }) => {
+        ok(typeof firstName === "string");
+        ok(typeof lastName === "string");
+        const fullName = firstName + lastName;
+        ok(fullName.toLowerCase().includes("alex"));
+        ok(fullName.toLowerCase().includes("yang"));
+        return uniqueId;
+      },
+      {
+        name: "unique_id",
+        description: "show user unique id",
+        parameters: {
+          type: "object",
+          properties: {
+            firstName: { type: "string" },
+            lastName: { type: "string" },
+          },
+          required: ["firstName", "lastName"],
+        },
+      },
+    );
+    const agent = new AnthropicAgent({
+      tools: [showUniqueId],
+    });
+    const { response } = await agent.chat({
+      message: "My name is Alex Yang. What is my unique id?",
+    });
+    consola.debug("response:", response);
+    ok(response.includes(uniqueId));
+  });
+
+  await t.test("sum numbers", async () => {
+    const openaiAgent = new AnthropicAgent({
+      tools: [sumNumbersTool],
+    });
+
+    const response = await openaiAgent.chat({
+      message: "how much is 1 + 1?",
+    });
+
+    ok(response.response.includes("2"));
+  });
+});

packages/core/e2e/node/fixtures/tools.ts

@@ -0,0 +1,47 @@
+import { FunctionTool } from "llamaindex";
+
+function sumNumbers({ a, b }: { a: number; b: number }) {
+  return `${a + b}`;
+}
+
+function divideNumbers({ a, b }: { a: number; b: number }) {
+  return `${a / b}`;
+}
+
+export const sumNumbersTool = FunctionTool.from(sumNumbers, {
+  name: "sumNumbers",
+  description: "Use this function to sum two numbers",
+  parameters: {
+    type: "object",
+    properties: {
+      a: {
+        type: "number",
+        description: "The first number",
+      },
+      b: {
+        type: "number",
+        description: "The second number",
+      },
+    },
+    required: ["a", "b"],
+  },
+});
+
+export const divideNumbersTool = FunctionTool.from(divideNumbers, {
+  name: "divideNumbers",
+  description: "Use this function to divide two numbers",
+  parameters: {
+    type: "object",
+    properties: {
+      a: {
+        type: "number",
+        description: "The first number",
+      },
+      b: {
+        type: "number",
+        description: "The second number",
+      },
+    },
+    required: ["a", "b"],
+  },
+});

packages/core/e2e/node/openai.e2e.ts

@@ -0,0 +1,233 @@
+import { consola } from "consola";
+import {
+  Document,
+  FunctionTool,
+  OpenAI,
+  OpenAIAgent,
+  QueryEngineTool,
+  Settings,
+  SubQuestionQueryEngine,
+  VectorStoreIndex,
+  type LLM,
+} from "llamaindex";
+import { ok, strictEqual } from "node:assert";
+import { beforeEach, test } from "node:test";
+import { divideNumbersTool, sumNumbersTool } from "./fixtures/tools.js";
+import { mockLLMEvent } from "./utils.js";
+
+let llm: LLM;
+beforeEach(async () => {
+  Settings.llm = new OpenAI({
+    model: "gpt-3.5-turbo",
+  });
+  llm = Settings.llm;
+});
+
+await test("openai llm", async (t) => {
+  await mockLLMEvent(t, "llm");
+  await t.test("llm.chat", async () => {
+    const response = await llm.chat({
+      messages: [
+        {
+          content: "Hello",
+          role: "user",
+        },
+      ],
+    });
+    consola.debug("response:", response);
+    ok(typeof response.message.content === "string");
+  });
+
+  await t.test("stream llm.chat", async () => {
+    const iter = await llm.chat({
+      stream: true,
+      messages: [
+        {
+          content: "hello",
+          role: "user",
+        },
+      ],
+    });
+    for await (const chunk of iter) {
+      consola.debug("chunk:", chunk);
+      ok(typeof chunk.delta === "string");
+    }
+  });
+});
+
+await test("gpt-4-turbo", async (t) => {
+  const llm = new OpenAI({ model: "gpt-4-turbo" });
+  Settings.llm = llm;
+  await mockLLMEvent(t, "gpt-4-turbo");
+  await t.test("agent", async () => {
+    const agent = new OpenAIAgent({
+      llm,
+      tools: [
+        {
+          call: async () => {
+            return "45 degrees and sunny in San Jose";
+          },
+          metadata: {
+            name: "Weather",
+            description: "Get the weather",
+            parameters: {
+              type: "object",
+              properties: {
+                location: { type: "string" },
+              },
+              required: ["location"],
+            },
+          },
+        },
+      ],
+    });
+    const { response } = await agent.chat({
+      message: "What is the weather in San Jose?",
+    });
+    consola.debug("response:", response);
+    ok(typeof response === "string");
+    ok(response.includes("45"));
+  });
+});
+
+await test("agent", async (t) => {
+  await mockLLMEvent(t, "agent");
+  await t.test("chat", async () => {
+    const agent = new OpenAIAgent({
+      tools: [
+        {
+          call: async () => {
+            return "35 degrees and sunny in San Francisco";
+          },
+          metadata: {
+            name: "Weather",
+            description: "Get the weather",
+            parameters: {
+              type: "object",
+              properties: {
+                location: { type: "string" },
+              },
+              required: ["location"],
+            },
+          },
+        },
+      ],
+    });
+    const result = await agent.chat({
+      message: "What is the weather in San Francisco?",
+    });
+    consola.debug("response:", result.response);
+    ok(typeof result.response === "string");
+    ok(result.response.includes("35"));
+  });
+
+  await t.test("async function", async () => {
+    const uniqueId = "123456789";
+    const showUniqueId = FunctionTool.from<{
+      firstName: string;
+      lastName: string;
+    }>(
+      async ({ firstName, lastName }) => {
+        ok(typeof firstName === "string");
+        ok(typeof lastName === "string");
+        const fullName = firstName + lastName;
+        ok(fullName.toLowerCase().includes("alex"));
+        ok(fullName.toLowerCase().includes("yang"));
+        return uniqueId;
+      },
+      {
+        name: "unique_id",
+        description: "show user unique id",
+        parameters: {
+          type: "object",
+          properties: {
+            firstName: { type: "string" },
+            lastName: { type: "string" },
+          },
+          required: ["firstName", "lastName"],
+        },
+      },
+    );
+    const agent = new OpenAIAgent({
+      tools: [showUniqueId],
+    });
+    const { response } = await agent.chat({
+      message: "My name is Alex Yang. What is my unique id?",
+    });
+    consola.debug("response:", response);
+    ok(response.includes(uniqueId));
+  });
+
+  await t.test("sum numbers", async () => {
+    const openaiAgent = new OpenAIAgent({
+      tools: [sumNumbersTool],
+    });
+
+    const response = await openaiAgent.chat({
+      message: "how much is 1 + 1?",
+    });
+
+    ok(response.response.includes("2"));
+  });
+});
+
+await test("agent stream", async (t) => {
+  await mockLLMEvent(t, "agent_stream");
+  await t.test("sum numbers stream", async (t) => {
+    const fn = t.mock.fn(() => {});
+    Settings.callbackManager.on("llm-tool-call", fn);
+
+    const agent = new OpenAIAgent({
+      tools: [sumNumbersTool, divideNumbersTool],
+    });
+
+    const { response } = await agent.chat({
+      message: "Divide 16 by 2 then add 20",
+      stream: true,
+    });
+
+    let message = "";
+
+    for await (const chunk of response) {
+      message += chunk.response;
+    }
+
+    strictEqual(fn.mock.callCount(), 2);
+    ok(message.includes("28"));
+    Settings.callbackManager.off("llm-tool-call", fn);
+  });
+});
+
+await test("queryEngine", async (t) => {
+  await mockLLMEvent(t, "queryEngine_subquestion");
+  await t.test("subquestion", async () => {
+    const fn = t.mock.fn(() => {});
+    Settings.callbackManager.on("llm-tool-call", fn);
+    const document = new Document({
+      text: "Bill Gates stole from Apple.\n Steve Jobs stole from Xerox.",
+    });
+    const index = await VectorStoreIndex.fromDocuments([document]);
+
+    const queryEngineTools = [
+      new QueryEngineTool({
+        queryEngine: index.asQueryEngine(),
+        metadata: {
+          name: "bill_gates_idea",
+          description: "Get what Bill Gates idea from.",
+        },
+      }),
+    ];
+
+    const queryEngine = SubQuestionQueryEngine.fromDefaults({
+      queryEngineTools,
+    });
+
+    const { response } = await queryEngine.query({
+      query: "What did Bill Gates steal from?",
+    });
+
+    ok(response.includes("Apple"));
+    strictEqual(fn.mock.callCount(), 0);
+    Settings.callbackManager.off("llm-tool-call", fn);
+  });
+});

packages/core/e2e/node/snapshot/agent.snap

@@ -0,0 +1,370 @@
+{
+  "llmEventStart": [
+    {
+      "id": "PRESERVE_0",
+      "messages": [
+        {
+          "content": "What is the weather in San Francisco?",
+          "role": "user"
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_1",
+      "messages": [
+        {
+          "content": "What is the weather in San Francisco?",
+          "role": "user"
+        },
+        {
+          "content": "",
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "Weather",
+              "input": "{\"location\":\"San Francisco\"}"
+            }
+          }
+        },
+        {
+          "content": "35 degrees and sunny in San Francisco",
+          "role": "user",
+          "options": {
+            "toolResult": {
+              "id": "HIDDEN",
+              "isError": false
+            }
+          }
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_2",
+      "messages": [
+        {
+          "content": "My name is Alex Yang. What is my unique id?",
+          "role": "user"
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_3",
+      "messages": [
+        {
+          "content": "My name is Alex Yang. What is my unique id?",
+          "role": "user"
+        },
+        {
+          "content": "",
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "unique_id",
+              "input": "{\"firstName\":\"Alex\",\"lastName\":\"Yang\"}"
+            }
+          }
+        },
+        {
+          "content": "123456789",
+          "role": "user",
+          "options": {
+            "toolResult": {
+              "id": "HIDDEN",
+              "isError": false
+            }
+          }
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_4",
+      "messages": [
+        {
+          "content": "how much is 1 + 1?",
+          "role": "user"
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_5",
+      "messages": [
+        {
+          "content": "how much is 1 + 1?",
+          "role": "user"
+        },
+        {
+          "content": "",
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "sumNumbers",
+              "input": "{\"a\":1,\"b\":1}"
+            }
+          }
+        },
+        {
+          "content": "2",
+          "role": "user",
+          "options": {
+            "toolResult": {
+              "id": "HIDDEN",
+              "isError": false
+            }
+          }
+        }
+      ]
+    }
+  ],
+  "llmEventEnd": [
+    {
+      "id": "PRESERVE_0",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": null,
+                "tool_calls": [
+                  {
+                    "id": "HIDDEN",
+                    "type": "function",
+                    "function": {
+                      "name": "Weather",
+                      "arguments": "{\"location\":\"San Francisco\"}"
+                    }
+                  }
+                ]
+              },
+              "logprobs": null,
+              "finish_reason": "tool_calls"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 49,
+            "completion_tokens": 15,
+            "total_tokens": 64
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "",
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "Weather",
+              "input": "{\"location\":\"San Francisco\"}"
+            }
+          }
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": "The weather in San Francisco is currently 35 degrees and sunny."
+              },
+              "logprobs": null,
+              "finish_reason": "stop"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 78,
+            "completion_tokens": 14,
+            "total_tokens": 92
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "The weather in San Francisco is currently 35 degrees and sunny.",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_2",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": null,
+                "tool_calls": [
+                  {
+                    "id": "HIDDEN",
+                    "type": "function",
+                    "function": {
+                      "name": "unique_id",
+                      "arguments": "{\"firstName\":\"Alex\",\"lastName\":\"Yang\"}"
+                    }
+                  }
+                ]
+              },
+              "logprobs": null,
+              "finish_reason": "tool_calls"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 59,
+            "completion_tokens": 18,
+            "total_tokens": 77
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "",
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "unique_id",
+              "input": "{\"firstName\":\"Alex\",\"lastName\":\"Yang\"}"
+            }
+          }
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_3",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": "Your unique id is 123456789."
+              },
+              "logprobs": null,
+              "finish_reason": "stop"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 88,
+            "completion_tokens": 10,
+            "total_tokens": 98
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "Your unique id is 123456789.",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_4",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": null,
+                "tool_calls": [
+                  {
+                    "id": "HIDDEN",
+                    "type": "function",
+                    "function": {
+                      "name": "sumNumbers",
+                      "arguments": "{\"a\":1,\"b\":1}"
+                    }
+                  }
+                ]
+              },
+              "logprobs": null,
+              "finish_reason": "tool_calls"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 70,
+            "completion_tokens": 18,
+            "total_tokens": 88
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "",
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "sumNumbers",
+              "input": "{\"a\":1,\"b\":1}"
+            }
+          }
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_5",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": "1 + 1 is equal to 2."
+              },
+              "logprobs": null,
+              "finish_reason": "stop"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 97,
+            "completion_tokens": 11,
+            "total_tokens": 108
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "1 + 1 is equal to 2.",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    }
+  ],
+  "llmEventStream": []
+}

packages/core/e2e/node/snapshot/anthropic-agent.snap

@@ -0,0 +1,403 @@
+{
+  "llmEventStart": [
+    {
+      "id": "PRESERVE_0",
+      "messages": [
+        {
+          "role": "user",
+          "content": "What is the weather in San Francisco?",
+          "options": {}
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_1",
+      "messages": [
+        {
+          "role": "user",
+          "content": "What is the weather in San Francisco?",
+          "options": {}
+        },
+        {
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe user is asking for the weather in a specific location, San Francisco. The Weather function is the relevant tool to answer this request, as it returns weather information for a given location.\n\nThe Weather function has one required parameter:\n- location (string): The user has directly provided the location of \"San Francisco\"\n\nSince the required location parameter has been provided by the user, we have all the necessary information to call the Weather function.\n</thinking>"
+            }
+          ],
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "Weather",
+              "input": {
+                "location": "San Francisco"
+              }
+            }
+          }
+        },
+        {
+          "content": "35 degrees and sunny in San Francisco",
+          "role": "user",
+          "options": {
+            "toolResult": {
+              "isError": false,
+              "id": "HIDDEN"
+            }
+          }
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_2",
+      "messages": [
+        {
+          "role": "user",
+          "content": "My name is Alex Yang. What is my unique id?",
+          "options": {}
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_3",
+      "messages": [
+        {
+          "role": "user",
+          "content": "My name is Alex Yang. What is my unique id?",
+          "options": {}
+        },
+        {
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe unique_id function is the relevant tool to answer the user's request for their unique ID. It requires two parameters:\nfirstName: The user provided their first name, which is \"Alex\"\nlastName: The user also provided their last name, \"Yang\"\nSince the user has provided all the required parameters, we can proceed with calling the unique_id function.\n</thinking>"
+            }
+          ],
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "unique_id",
+              "input": {
+                "firstName": "Alex",
+                "lastName": "Yang"
+              }
+            }
+          }
+        },
+        {
+          "content": "123456789",
+          "role": "user",
+          "options": {
+            "toolResult": {
+              "isError": false,
+              "id": "HIDDEN"
+            }
+          }
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_4",
+      "messages": [
+        {
+          "role": "user",
+          "content": "how much is 1 + 1?",
+          "options": {}
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_5",
+      "messages": [
+        {
+          "role": "user",
+          "content": "how much is 1 + 1?",
+          "options": {}
+        },
+        {
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe user is asking to sum the numbers 1 and 1. The relevant tool to use is the sumNumbers function, which takes two number parameters a and b.\nThe user has directly provided the values for the parameters:\na = 1 \nb = 1\nSince all the required parameters have been provided, we can proceed with calling the function.\n</thinking>"
+            }
+          ],
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "sumNumbers",
+              "input": {
+                "a": 1,
+                "b": 1
+              }
+            }
+          }
+        },
+        {
+          "content": "2",
+          "role": "user",
+          "options": {
+            "toolResult": {
+              "isError": false,
+              "id": "HIDDEN"
+            }
+          }
+        }
+      ]
+    }
+  ],
+  "llmEventEnd": [
+    {
+      "id": "PRESERVE_0",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "type": "message",
+          "role": "assistant",
+          "model": "claude-3-opus-20240229",
+          "stop_sequence": null,
+          "usage": {
+            "input_tokens": 462,
+            "output_tokens": 147
+          },
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe user is asking for the weather in a specific location, San Francisco. The Weather function is the relevant tool to answer this request, as it returns weather information for a given location.\n\nThe Weather function has one required parameter:\n- location (string): The user has directly provided the location of \"San Francisco\"\n\nSince the required location parameter has been provided by the user, we have all the necessary information to call the Weather function.\n</thinking>"
+            },
+            {
+              "type": "tool_use",
+              "id": "HIDDEN",
+              "name": "Weather",
+              "input": {
+                "location": "San Francisco"
+              }
+            }
+          ],
+          "stop_reason": "tool_use"
+        },
+        "message": {
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe user is asking for the weather in a specific location, San Francisco. The Weather function is the relevant tool to answer this request, as it returns weather information for a given location.\n\nThe Weather function has one required parameter:\n- location (string): The user has directly provided the location of \"San Francisco\"\n\nSince the required location parameter has been provided by the user, we have all the necessary information to call the Weather function.\n</thinking>"
+            }
+          ],
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "Weather",
+              "input": {
+                "location": "San Francisco"
+              }
+            }
+          }
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "type": "message",
+          "role": "assistant",
+          "model": "claude-3-opus-20240229",
+          "stop_sequence": null,
+          "usage": {
+            "input_tokens": 628,
+            "output_tokens": 18
+          },
+          "content": [
+            {
+              "type": "text",
+              "text": "The current weather in San Francisco is 35 degrees and sunny."
+            }
+          ],
+          "stop_reason": "end_turn"
+        },
+        "message": {
+          "content": [
+            {
+              "type": "text",
+              "text": "The current weather in San Francisco is 35 degrees and sunny."
+            }
+          ],
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_2",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "type": "message",
+          "role": "assistant",
+          "model": "claude-3-opus-20240229",
+          "stop_sequence": null,
+          "usage": {
+            "input_tokens": 482,
+            "output_tokens": 152
+          },
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe unique_id function is the relevant tool to answer the user's request for their unique ID. It requires two parameters:\nfirstName: The user provided their first name, which is \"Alex\"\nlastName: The user also provided their last name, \"Yang\"\nSince the user has provided all the required parameters, we can proceed with calling the unique_id function.\n</thinking>"
+            },
+            {
+              "type": "tool_use",
+              "id": "HIDDEN",
+              "name": "unique_id",
+              "input": {
+                "firstName": "Alex",
+                "lastName": "Yang"
+              }
+            }
+          ],
+          "stop_reason": "tool_use"
+        },
+        "message": {
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe unique_id function is the relevant tool to answer the user's request for their unique ID. It requires two parameters:\nfirstName: The user provided their first name, which is \"Alex\"\nlastName: The user also provided their last name, \"Yang\"\nSince the user has provided all the required parameters, we can proceed with calling the unique_id function.\n</thinking>"
+            }
+          ],
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "unique_id",
+              "input": {
+                "firstName": "Alex",
+                "lastName": "Yang"
+              }
+            }
+          }
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_3",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "type": "message",
+          "role": "assistant",
+          "model": "claude-3-opus-20240229",
+          "stop_sequence": null,
+          "usage": {
+            "input_tokens": 648,
+            "output_tokens": 13
+          },
+          "content": [
+            {
+              "type": "text",
+              "text": "Your unique ID is 123456789."
+            }
+          ],
+          "stop_reason": "end_turn"
+        },
+        "message": {
+          "content": [
+            {
+              "type": "text",
+              "text": "Your unique ID is 123456789."
+            }
+          ],
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_4",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "type": "message",
+          "role": "assistant",
+          "model": "claude-3-opus-20240229",
+          "stop_sequence": null,
+          "usage": {
+            "input_tokens": 498,
+            "output_tokens": 151
+          },
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe user is asking to sum the numbers 1 and 1. The relevant tool to use is the sumNumbers function, which takes two number parameters a and b.\nThe user has directly provided the values for the parameters:\na = 1 \nb = 1\nSince all the required parameters have been provided, we can proceed with calling the function.\n</thinking>"
+            },
+            {
+              "type": "tool_use",
+              "id": "HIDDEN",
+              "name": "sumNumbers",
+              "input": {
+                "a": 1,
+                "b": 1
+              }
+            }
+          ],
+          "stop_reason": "tool_use"
+        },
+        "message": {
+          "content": [
+            {
+              "type": "text",
+              "text": "<thinking>\nThe user is asking to sum the numbers 1 and 1. The relevant tool to use is the sumNumbers function, which takes two number parameters a and b.\nThe user has directly provided the values for the parameters:\na = 1 \nb = 1\nSince all the required parameters have been provided, we can proceed with calling the function.\n</thinking>"
+            }
+          ],
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "sumNumbers",
+              "input": {
+                "a": 1,
+                "b": 1
+              }
+            }
+          }
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_5",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "type": "message",
+          "role": "assistant",
+          "model": "claude-3-opus-20240229",
+          "stop_sequence": null,
+          "usage": {
+            "input_tokens": 661,
+            "output_tokens": 16
+          },
+          "content": [
+            {
+              "type": "text",
+              "text": "So 1 + 1 = 2."
+            }
+          ],
+          "stop_reason": "end_turn"
+        },
+        "message": {
+          "content": [
+            {
+              "type": "text",
+              "text": "So 1 + 1 = 2."
+            }
+          ],
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    }
+  ],
+  "llmEventStream": []
+}

packages/core/e2e/node/snapshot/gpt-4-turbo.snap

@@ -0,0 +1,128 @@
+{
+  "llmEventStart": [
+    {
+      "id": "PRESERVE_0",
+      "messages": [
+        {
+          "content": "What is the weather in San Jose?",
+          "role": "user"
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_1",
+      "messages": [
+        {
+          "content": "What is the weather in San Jose?",
+          "role": "user"
+        },
+        {
+          "content": "",
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "Weather",
+              "input": "{\"location\":\"San Jose\"}"
+            }
+          }
+        },
+        {
+          "content": "45 degrees and sunny in San Jose",
+          "role": "user",
+          "options": {
+            "toolResult": {
+              "id": "HIDDEN",
+              "isError": false
+            }
+          }
+        }
+      ]
+    }
+  ],
+  "llmEventEnd": [
+    {
+      "id": "PRESERVE_0",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": null,
+                "tool_calls": [
+                  {
+                    "id": "HIDDEN",
+                    "type": "function",
+                    "function": {
+                      "name": "Weather",
+                      "arguments": "{\"location\":\"San Jose\"}"
+                    }
+                  }
+                ]
+              },
+              "logprobs": null,
+              "finish_reason": "tool_calls"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 49,
+            "completion_tokens": 15,
+            "total_tokens": 64
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "",
+          "role": "assistant",
+          "options": {
+            "toolCall": {
+              "id": "HIDDEN",
+              "name": "Weather",
+              "input": "{\"location\":\"San Jose\"}"
+            }
+          }
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": "The weather in San Jose is currently 45 degrees and sunny."
+              },
+              "logprobs": null,
+              "finish_reason": "stop"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 78,
+            "completion_tokens": 14,
+            "total_tokens": 92
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "The weather in San Jose is currently 45 degrees and sunny.",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    }
+  ],
+  "llmEventStream": []
+}

packages/core/e2e/node/snapshot/llm-anthropic.snap

@@ -0,0 +1,310 @@
+{
+  "llmEventStart": [
+    {
+      "id": "PRESERVE_0",
+      "messages": [
+        {
+          "content": "Hello",
+          "role": "user",
+          "options": {}
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_1",
+      "messages": [
+        {
+          "content": "hello",
+          "role": "user"
+        }
+      ]
+    }
+  ],
+  "llmEventEnd": [
+    {
+      "id": "PRESERVE_0",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "type": "message",
+          "role": "assistant",
+          "model": "claude-3-opus-20240229",
+          "stop_sequence": null,
+          "usage": {
+            "input_tokens": 8,
+            "output_tokens": 12
+          },
+          "content": [
+            {
+              "type": "text",
+              "text": "Hello! How can I assist you today?"
+            }
+          ],
+          "stop_reason": "end_turn"
+        },
+        "message": {
+          "content": "Hello! How can I assist you today?",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "response": {
+        "raw": [
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": "Hello"
+              }
+            },
+            "delta": "Hello",
+            "options": {}
+          },
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": "!"
+              }
+            },
+            "delta": "!",
+            "options": {}
+          },
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": " How"
+              }
+            },
+            "delta": " How",
+            "options": {}
+          },
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": " can"
+              }
+            },
+            "delta": " can",
+            "options": {}
+          },
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": " I"
+              }
+            },
+            "delta": " I",
+            "options": {}
+          },
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": " assist"
+              }
+            },
+            "delta": " assist",
+            "options": {}
+          },
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": " you"
+              }
+            },
+            "delta": " you",
+            "options": {}
+          },
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": " today"
+              }
+            },
+            "delta": " today",
+            "options": {}
+          },
+          {
+            "raw": {
+              "type": "content_block_delta",
+              "index": 0,
+              "delta": {
+                "type": "text_delta",
+                "text": "?"
+              }
+            },
+            "delta": "?",
+            "options": {}
+          }
+        ],
+        "message": {
+          "content": "Hello! How can I assist you today?",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    }
+  ],
+  "llmEventStream": [
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": "Hello"
+          }
+        },
+        "delta": "Hello",
+        "options": {}
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": "!"
+          }
+        },
+        "delta": "!",
+        "options": {}
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": " How"
+          }
+        },
+        "delta": " How",
+        "options": {}
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": " can"
+          }
+        },
+        "delta": " can",
+        "options": {}
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": " I"
+          }
+        },
+        "delta": " I",
+        "options": {}
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": " assist"
+          }
+        },
+        "delta": " assist",
+        "options": {}
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": " you"
+          }
+        },
+        "delta": " you",
+        "options": {}
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": " today"
+          }
+        },
+        "delta": " today",
+        "options": {}
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "type": "content_block_delta",
+          "index": 0,
+          "delta": {
+            "type": "text_delta",
+            "text": "?"
+          }
+        },
+        "delta": "?",
+        "options": {}
+      }
+    }
+  ]
+}

packages/core/e2e/node/snapshot/llm.snap

@@ -0,0 +1,476 @@
+{
+  "llmEventStart": [
+    {
+      "id": "PRESERVE_0",
+      "messages": [
+        {
+          "content": "Hello",
+          "role": "user"
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_1",
+      "messages": [
+        {
+          "content": "hello",
+          "role": "user"
+        }
+      ]
+    }
+  ],
+  "llmEventEnd": [
+    {
+      "id": "PRESERVE_0",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": "Hello! How can I assist you today?"
+              },
+              "logprobs": null,
+              "finish_reason": "stop"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 8,
+            "completion_tokens": 9,
+            "total_tokens": 17
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "Hello! How can I assist you today?",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "response": {
+        "raw": [
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": "Hello"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": "Hello"
+          },
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": "!"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": "!"
+          },
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": " How"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": " How"
+          },
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": " can"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": " can"
+          },
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": " I"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": " I"
+          },
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": " assist"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": " assist"
+          },
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": " you"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": " you"
+          },
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": " today"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": " today"
+          },
+          {
+            "raw": {
+              "id": "HIDDEN",
+              "object": "chat.completion.chunk",
+              "created": 114514,
+              "model": "gpt-3.5-turbo-0125",
+              "system_fingerprint": "HIDDEN",
+              "choices": [
+                {
+                  "index": 0,
+                  "delta": {
+                    "content": "?"
+                  },
+                  "logprobs": null,
+                  "finish_reason": null
+                }
+              ]
+            },
+            "options": {},
+            "delta": "?"
+          }
+        ],
+        "message": {
+          "content": "Hello! How can I assist you today?",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    }
+  ],
+  "llmEventStream": [
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": "Hello"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": "Hello"
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": "!"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": "!"
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": " How"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": " How"
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": " can"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": " can"
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": " I"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": " I"
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": " assist"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": " assist"
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": " you"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": " you"
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": " today"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": " today"
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "chunk": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion.chunk",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "system_fingerprint": "HIDDEN",
+          "choices": [
+            {
+              "index": 0,
+              "delta": {
+                "content": "?"
+              },
+              "logprobs": null,
+              "finish_reason": null
+            }
+          ]
+        },
+        "options": {},
+        "delta": "?"
+      }
+    }
+  ]
+}

packages/core/e2e/node/snapshot/queryEngine_subquestion.snap

@@ -0,0 +1,133 @@
+{
+  "llmEventStart": [
+    {
+      "id": "PRESERVE_0",
+      "messages": [
+        {
+          "content": "Given a user question, and a list of tools, output a list of relevant sub-questions that when composed can help answer the full user question:\n\n# Example 1\n<Tools>\n```json\n{\n    \"uber_10k\": \"Provides information about Uber financials for year 2021\",\n    \"lyft_10k\": \"Provides information about Lyft financials for year 2021\"\n}\n```\n\n<User Question>\nCompare and contrast the revenue growth and EBITDA of Uber and Lyft for year 2021\n\n<Output>\n```json\n[\n    {\n        \"subQuestion\": \"What is the revenue growth of Uber\",\n        \"toolName\": \"uber_10k\"\n    },\n    {\n        \"subQuestion\": \"What is the EBITDA of Uber\",\n        \"toolName\": \"uber_10k\"\n    },\n    {\n        \"subQuestion\": \"What is the revenue growth of Lyft\",\n        \"toolName\": \"lyft_10k\"\n    },\n    {\n        \"subQuestion\": \"What is the EBITDA of Lyft\",\n        \"toolName\": \"lyft_10k\"\n    }\n]\n```\n\n# Example 2\n<Tools>\n```json\n{\n    \"bill_gates_idea\": \"Get what Bill Gates idea from.\"\n}\n```\n\n<User Question>\nWhat did Bill Gates steal from?\n\n<Output>\n",
+          "role": "user"
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_1",
+      "messages": [
+        {
+          "content": "Context information is below.\n---------------------\nBill Gates stole from Apple. Steve Jobs stole from Xerox.\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: What is Bill Gates' idea\nAnswer:",
+          "role": "user"
+        }
+      ]
+    },
+    {
+      "id": "PRESERVE_2",
+      "messages": [
+        {
+          "content": "Context information is below.\n---------------------\nSub question: What is Bill Gates' idea\nResponse: Bill Gates' idea was to steal from Apple.\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: What did Bill Gates steal from?\nAnswer:",
+          "role": "user"
+        }
+      ]
+    }
+  ],
+  "llmEventEnd": [
+    {
+      "id": "PRESERVE_0",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": "```json\n[\n    {\n        \"subQuestion\": \"What is Bill Gates' idea\",\n        \"toolName\": \"bill_gates_idea\"\n    }\n]\n```"
+              },
+              "logprobs": null,
+              "finish_reason": "stop"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 290,
+            "completion_tokens": 35,
+            "total_tokens": 325
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "```json\n[\n    {\n        \"subQuestion\": \"What is Bill Gates' idea\",\n        \"toolName\": \"bill_gates_idea\"\n    }\n]\n```",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_1",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": "Bill Gates' idea was to steal from Apple."
+              },
+              "logprobs": null,
+              "finish_reason": "stop"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 53,
+            "completion_tokens": 10,
+            "total_tokens": 63
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "Bill Gates' idea was to steal from Apple.",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    },
+    {
+      "id": "PRESERVE_2",
+      "response": {
+        "raw": {
+          "id": "HIDDEN",
+          "object": "chat.completion",
+          "created": 114514,
+          "model": "gpt-3.5-turbo-0125",
+          "choices": [
+            {
+              "index": 0,
+              "message": {
+                "role": "assistant",
+                "content": "Bill Gates stole from Apple."
+              },
+              "logprobs": null,
+              "finish_reason": "stop"
+            }
+          ],
+          "usage": {
+            "prompt_tokens": 62,
+            "completion_tokens": 6,
+            "total_tokens": 68
+          },
+          "system_fingerprint": "HIDDEN"
+        },
+        "message": {
+          "content": "Bill Gates stole from Apple.",
+          "role": "assistant",
+          "options": {}
+        }
+      }
+    }
+  ],
+  "llmEventStream": []
+}

packages/core/e2e/node/utils.ts

@@ -0,0 +1,139 @@
+/* eslint-disable turbo/no-undeclared-env-vars */
+import {
+  Settings,
+  type LLMEndEvent,
+  type LLMStartEvent,
+  type LLMStreamEvent,
+} from "llamaindex";
+import { readFile, writeFile } from "node:fs/promises";
+import { join } from "node:path";
+import { type test } from "node:test";
+import { fileURLToPath } from "node:url";
+
+type MockStorage = {
+  llmEventStart: LLMStartEvent["detail"]["payload"][];
+  llmEventEnd: LLMEndEvent["detail"]["payload"][];
+  llmEventStream: LLMStreamEvent["detail"]["payload"][];
+};
+
+export const llmCompleteMockStorage: MockStorage = {
+  llmEventStart: [],
+  llmEventEnd: [],
+  llmEventStream: [],
+};
+
+export const testRootDir = fileURLToPath(new URL(".", import.meta.url));
+
+export async function mockLLMEvent(
+  t: Parameters<NonNullable<Parameters<typeof test>[0]>>[0],
+  snapshotName: string,
+) {
+  const idMap = new Map<string, string>();
+  let counter = 0;
+  const newLLMCompleteMockStorage: MockStorage = {
+    llmEventStart: [],
+    llmEventEnd: [],
+    llmEventStream: [],
+  };
+
+  function captureLLMStart(event: LLMStartEvent) {
+    idMap.set(event.detail.payload.id, `PRESERVE_${counter++}`);
+    newLLMCompleteMockStorage.llmEventStart.push({
+      ...event.detail.payload,
+      // @ts-expect-error id is not UUID, but it is fine for testing
+      id: idMap.get(event.detail.payload.id)!,
+    });
+  }
+
+  function captureLLMEnd(event: LLMEndEvent) {
+    newLLMCompleteMockStorage.llmEventEnd.push({
+      ...event.detail.payload,
+      // @ts-expect-error id is not UUID, but it is fine for testing
+      id: idMap.get(event.detail.payload.id)!,
+    });
+  }
+
+  function captureLLMStream(event: LLMStreamEvent) {
+    newLLMCompleteMockStorage.llmEventStream.push({
+      ...event.detail.payload,
+      // @ts-expect-error id is not UUID, but it is fine for testing
+      id: idMap.get(event.detail.payload.id)!,
+    });
+  }
+
+  await readFile(join(testRootDir, "snapshot", `${snapshotName}.snap`), {
+    encoding: "utf-8",
+  })
+    .then((data) => {
+      const result = JSON.parse(data) as MockStorage;
+      result["llmEventEnd"].forEach((event) => {
+        llmCompleteMockStorage.llmEventEnd.push(event);
+      });
+      result["llmEventStart"].forEach((event) => {
+        llmCompleteMockStorage.llmEventStart.push(event);
+      });
+      result["llmEventStream"].forEach((event) => {
+        llmCompleteMockStorage.llmEventStream.push(event);
+      });
+    })
+    .catch((error) => {
+      if (error.code === "ENOENT") {
+        console.warn("Snapshot file not found, will create a new one");
+        return;
+      }
+    });
+  Settings.callbackManager.on("llm-start", captureLLMStart);
+  Settings.callbackManager.on("llm-end", captureLLMEnd);
+  Settings.callbackManager.on("llm-stream", captureLLMStream);
+
+  t.after(async () => {
+    Settings.callbackManager.off("llm-stream", captureLLMStream);
+    Settings.callbackManager.off("llm-end", captureLLMEnd);
+    Settings.callbackManager.off("llm-start", captureLLMStart);
+    // eslint-disable-next-line turbo/no-undeclared-env-vars
+    if (process.env.UPDATE_SNAPSHOT === "1") {
+      const data = JSON.stringify(newLLMCompleteMockStorage, null, 2)
+        .replace(/"id": "(?!PRESERVE_).*"/g, '"id": "HIDDEN"')
+        .replace(/"created": \d+/g, `"created": 114514`)
+        .replace(
+          /"system_fingerprint": ".*"/g,
+          '"system_fingerprint": "HIDDEN"',
+        )
+        .replace(/"tool_call_id": ".*"/g, '"tool_call_id": "HIDDEN"');
+      await writeFile(
+        join(testRootDir, "snapshot", `${snapshotName}.snap`),
+        data,
+      );
+      return;
+    }
+    if (
+      newLLMCompleteMockStorage.llmEventEnd.length !==
+      llmCompleteMockStorage.llmEventEnd.length
+    ) {
+      throw new Error("New LLMEndEvent does not match, please update snapshot");
+    }
+    if (
+      newLLMCompleteMockStorage.llmEventStart.length !==
+      llmCompleteMockStorage.llmEventStart.length
+    ) {
+      throw new Error(
+        "New LLMStartEvent does not match, please update snapshot",
+      );
+    }
+
+    if (
+      newLLMCompleteMockStorage.llmEventStream.length !==
+      llmCompleteMockStorage.llmEventStream.length
+    ) {
+      throw new Error(
+        "New LLMStreamEvent does not match, please update snapshot",
+      );
+    }
+  });
+  // cleanup
+  t.after(() => {
+    llmCompleteMockStorage.llmEventEnd = [];
+    llmCompleteMockStorage.llmEventStart = [];
+    llmCompleteMockStorage.llmEventStream = [];
+  });
+}

packages/core/e2e/package.json

@@ -5,7 +5,8 @@
   "type": "module",
   "scripts": {
     "e2e": "node --import tsx --import ./mock-register.js --test ./node/*.e2e.ts",
-    "e2e:nomock": "node --import tsx --test ./node/*.e2e.ts"
+    "e2e:nomock": "node --import tsx --test ./node/*.e2e.ts",
+    "e2e:updatesnap": "UPDATE_SNAPSHOT=1 node --import tsx --test ./node/*.e2e.ts"
   },
   "devDependencies": {
     "@faker-js/faker": "^8.4.1",

packages/core/package.json

@@ -1,29 +1,30 @@
 {
   "name": "llamaindex",
-  "version": "0.2.3",
+  "version": "0.2.9",
   "expectedMinorVersion": "2",
   "license": "MIT",
   "type": "module",
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.18.0",
+    "@anthropic-ai/sdk": "^0.20.4",
     "@aws-crypto/sha256-js": "^5.2.0",
     "@datastax/astra-db-ts": "^0.1.4",
     "@grpc/grpc-js": "^1.10.6",
     "@llamaindex/cloud": "0.0.5",
     "@llamaindex/env": "workspace:*",
     "@mistralai/mistralai": "^0.1.3",
-    "@notionhq/client": "^2.2.14",
+    "@notionhq/client": "^2.2.15",
     "@pinecone-database/pinecone": "^2.2.0",
     "@qdrant/js-client-rest": "^1.8.2",
     "@types/lodash": "^4.17.0",
-    "@types/node": "^18.19.31",
+    "@types/node": "^20.12.7",
     "@types/papaparse": "^5.3.14",
     "@types/pg": "^8.11.5",
     "@xenova/transformers": "^2.16.1",
     "@zilliz/milvus2-sdk-node": "^2.3.5",
+    "ajv": "^8.12.0",
     "assemblyai": "^4.3.4",
     "chromadb": "~1.7.3",
-    "cohere-ai": "^7.9.2",
+    "cohere-ai": "^7.9.3",
     "js-tiktoken": "^1.0.10",
     "lodash": "^4.17.21",
     "magic-bytes.js": "^1.10.0",
@@ -50,7 +51,7 @@
     "concurrently": "^8.2.2",
     "glob": "^10.3.12",
     "madge": "^6.1.0",
-    "typescript": "^5.4.4"
+    "typescript": "^5.4.5"
   },
   "engines": {
     "node": ">=18.0.0"

packages/core/src/ChatHistory.ts

@@ -8,20 +8,22 @@ import { extractText } from "./llm/utils.js";
 /**
  * A ChatHistory is used to keep the state of back and forth chat messages
  */
-export abstract class ChatHistory {
-  abstract get messages(): ChatMessage[];
+export abstract class ChatHistory<
+  AdditionalMessageOptions extends object = object,
+> {
+  abstract get messages(): ChatMessage<AdditionalMessageOptions>[];
   /**
    * Adds a message to the chat history.
    * @param message
    */
-  abstract addMessage(message: ChatMessage): void;
+  abstract addMessage(message: ChatMessage<AdditionalMessageOptions>): void;
 
   /**
    * Returns the messages that should be used as input to the LLM.
    */
   abstract requestMessages(
-    transientMessages?: ChatMessage[],
-  ): Promise<ChatMessage[]>;
+    transientMessages?: ChatMessage<AdditionalMessageOptions>[],
+  ): Promise<ChatMessage<AdditionalMessageOptions>[]>;
 
   /**
    * Resets the chat history so that it's empty.
@@ -31,7 +33,7 @@ export abstract class ChatHistory {
   /**
    * Returns the new messages since the last call to this function (or since calling the constructor)
    */
-  abstract newMessages(): ChatMessage[];
+  abstract newMessages(): ChatMessage<AdditionalMessageOptions>[];
 }
 
 export class SimpleChatHistory extends ChatHistory {
@@ -108,6 +110,7 @@ export class SummaryChatHistory extends ChatHistory {
             context: messagesToHistoryStr(messagesToSummarize),
           }),
           role: "user" as MessageType,
+          options: {},
         },
       ];
       // remove oldest message until the chat history is short enough for the context window
@@ -116,7 +119,9 @@ export class SummaryChatHistory extends ChatHistory {
       this.tokenizer(promptMessages[0].content).length > this.tokensToSummarize
     );
 
-    const response = await this.llm.chat({ messages: promptMessages });
+    const response = await this.llm.chat({
+      messages: promptMessages,
+    });
     return { content: response.message.content, role: "memory" };
   }
 

packages/core/src/Response.ts

@@ -1,14 +1,14 @@
-import type { BaseNode } from "./Node.js";
+import type { NodeWithScore } from "./Node.js";
 
 /**
  * Response is the output of a LLM
  */
 export class Response {
   response: string;
-  sourceNodes?: BaseNode[];
+  sourceNodes?: NodeWithScore[];
   metadata: Record<string, unknown> = {};
 
-  constructor(response: string, sourceNodes?: BaseNode[]) {
+  constructor(response: string, sourceNodes?: NodeWithScore[]) {
     this.response = response;
     this.sourceNodes = sourceNodes || [];
   }

packages/core/src/agent/anthropic.ts

@@ -0,0 +1,170 @@
+import { Settings } from "../Settings.js";
+import {
+  AgentChatResponse,
+  type ChatEngineParamsNonStreaming,
+} from "../engines/chat/index.js";
+import { wrapEventCaller } from "../internal/context/EventCaller.js";
+import { getCallbackManager } from "../internal/settings/CallbackManager.js";
+import { prettifyError } from "../internal/utils.js";
+import { Anthropic } from "../llm/anthropic.js";
+import type {
+  ChatMessage,
+  ChatResponse,
+  ToolCallLLMMessageOptions,
+} from "../llm/index.js";
+import { extractText } from "../llm/utils.js";
+import { ObjectRetriever } from "../objects/index.js";
+import type { BaseToolWithCall } from "../types.js";
+
+const MAX_TOOL_CALLS = 10;
+
+type AnthropicParamsBase = {
+  llm?: Anthropic;
+  chatHistory?: ChatMessage<ToolCallLLMMessageOptions>[];
+};
+
+type AnthropicParamsWithTools = AnthropicParamsBase & {
+  tools: BaseToolWithCall[];
+};
+
+type AnthropicParamsWithToolRetriever = AnthropicParamsBase & {
+  toolRetriever: ObjectRetriever<BaseToolWithCall>;
+};
+
+export type AnthropicAgentParams =
+  | AnthropicParamsWithTools
+  | AnthropicParamsWithToolRetriever;
+
+type AgentContext = {
+  toolCalls: number;
+  llm: Anthropic;
+  tools: BaseToolWithCall[];
+  messages: ChatMessage<ToolCallLLMMessageOptions>[];
+  shouldContinue: (context: AgentContext) => boolean;
+};
+
+type TaskResult = {
+  response: ChatResponse<ToolCallLLMMessageOptions>;
+  chatHistory: ChatMessage<ToolCallLLMMessageOptions>[];
+};
+
+async function task(
+  context: AgentContext,
+  input: ChatMessage<ToolCallLLMMessageOptions>,
+): Promise<TaskResult> {
+  const { llm, tools, messages } = context;
+  const nextMessages: ChatMessage<ToolCallLLMMessageOptions>[] = [
+    ...messages,
+    input,
+  ];
+  const response = await llm.chat({
+    stream: false,
+    tools,
+    messages: nextMessages,
+  });
+  const options = response.message.options ?? {};
+  if ("toolCall" in options) {
+    const { toolCall } = options;
+    const { input, name, id } = toolCall;
+    const targetTool = tools.find((tool) => tool.metadata.name === name);
+    let output: string;
+    let isError = true;
+    if (!context.shouldContinue(context)) {
+      output = "Error: Tool call limit reached";
+    } else if (!targetTool) {
+      output = `Error: Tool ${name} not found`;
+    } else {
+      try {
+        getCallbackManager().dispatchEvent("llm-tool-call", {
+          payload: {
+            toolCall: {
+              name,
+              input,
+            },
+          },
+        });
+        output = await targetTool.call(input);
+        isError = false;
+      } catch (error: unknown) {
+        output = prettifyError(error);
+      }
+    }
+    return task(
+      {
+        ...context,
+        toolCalls: context.toolCalls + 1,
+        messages: [...nextMessages, response.message],
+      },
+      {
+        content: output,
+        role: "user",
+        options: {
+          toolResult: {
+            isError,
+            id,
+          },
+        },
+      },
+    );
+  } else {
+    return { response, chatHistory: [...nextMessages, response.message] };
+  }
+}
+
+export class AnthropicAgent {
+  readonly #llm: Anthropic;
+  readonly #tools:
+    | BaseToolWithCall[]
+    | ((query: string) => Promise<BaseToolWithCall[]>) = [];
+  #chatHistory: ChatMessage<ToolCallLLMMessageOptions>[] = [];
+
+  constructor(params: AnthropicAgentParams) {
+    this.#llm =
+      params.llm ?? Settings.llm instanceof Anthropic
+        ? (Settings.llm as Anthropic)
+        : new Anthropic();
+    if ("tools" in params) {
+      this.#tools = params.tools;
+    } else if ("toolRetriever" in params) {
+      this.#tools = params.toolRetriever.retrieve.bind(params.toolRetriever);
+    }
+    if (Array.isArray(params.chatHistory)) {
+      this.#chatHistory = params.chatHistory;
+    }
+  }
+
+  static shouldContinue(context: AgentContext): boolean {
+    return context.toolCalls < MAX_TOOL_CALLS;
+  }
+
+  public reset(): void {
+    this.#chatHistory = [];
+  }
+
+  getTools(query: string): Promise<BaseToolWithCall[]> | BaseToolWithCall[] {
+    return typeof this.#tools === "function" ? this.#tools(query) : this.#tools;
+  }
+
+  @wrapEventCaller
+  async chat(
+    params: ChatEngineParamsNonStreaming,
+  ): Promise<Promise<AgentChatResponse>> {
+    const { chatHistory, response } = await task(
+      {
+        llm: this.#llm,
+        tools: await this.getTools(extractText(params.message)),
+        toolCalls: 0,
+        messages: [...this.#chatHistory],
+        // do we need this?
+        shouldContinue: AnthropicAgent.shouldContinue,
+      },
+      {
+        role: "user",
+        content: params.message,
+        options: {},
+      },
+    );
+    this.#chatHistory = [...chatHistory];
+    return new AgentChatResponse(extractText(response.message.content));
+  }
+}

packages/core/src/agent/index.ts

@@ -1,3 +1,5 @@
+// Not exporting the AnthropicAgent because it is not ready to ship yet.
+// export { AnthropicAgent, type AnthropicAgentParams } from "./anthropic.js";
 export * from "./openai/base.js";
 export * from "./openai/worker.js";
 export * from "./react/base.js";

packages/core/src/agent/openai/base.ts

@@ -1,6 +1,7 @@
 import { Settings } from "../../Settings.js";
 import type { ChatMessage } from "../../llm/index.js";
 import { OpenAI } from "../../llm/index.js";
+import type { BaseMemory } from "../../memory/types.js";
 import type { ObjectRetriever } from "../../objects/base.js";
 import type { BaseTool } from "../../types.js";
 import { AgentRunner } from "../runner/base.js";
@@ -9,11 +10,11 @@ import { OpenAIAgentWorker } from "./worker.js";
 type OpenAIAgentParams = {
   tools?: BaseTool[];
   llm?: OpenAI;
-  memory?: any;
+  memory?: BaseMemory;
   prefixMessages?: ChatMessage[];
   maxFunctionCalls?: number;
   defaultToolChoice?: string;
-  toolRetriever?: ObjectRetriever;
+  toolRetriever?: ObjectRetriever<BaseTool>;
   systemPrompt?: string;
 };
 
@@ -55,7 +56,7 @@ export class OpenAIAgent extends AgentRunner {
       ];
     }
 
-    if (!llm?.metadata.isFunctionCallingModel) {
+    if (!llm?.supportToolCall) {
       throw new Error("LLM model must be a function-calling model");
     }
 
@@ -72,6 +73,7 @@ export class OpenAIAgent extends AgentRunner {
       llm,
       memory,
       defaultToolChoice,
+      // @ts-expect-error 2322
       chatHistory: prefixMessages,
     });
   }

packages/core/src/agent/openai/worker.ts

@@ -1,4 +1,4 @@
-import { randomUUID } from "@llamaindex/env";
+import { pipeline, randomUUID } from "@llamaindex/env";
 import type { ChatCompletionToolChoiceOption } from "openai/resources/chat/completions";
 import { Response } from "../../Response.js";
 import { Settings } from "../../Settings.js";
@@ -14,12 +14,10 @@ import {
   type ChatResponseChunk,
   type LLMChatParamsBase,
   type OpenAIAdditionalChatOptions,
+  type ToolCallLLMMessageOptions,
+  type ToolCallOptions,
 } from "../../llm/index.js";
-import {
-  extractText,
-  streamConverter,
-  streamReducer,
-} from "../../llm/utils.js";
+import { extractText } from "../../llm/utils.js";
 import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer.js";
 import type { ObjectRetriever } from "../../objects/base.js";
 import type { ToolOutput } from "../../tools/types.js";
@@ -28,28 +26,25 @@ import type { BaseTool } from "../../types.js";
 import type { AgentWorker, Task } from "../types.js";
 import { TaskStep, TaskStepOutput } from "../types.js";
 import { addUserStepToMemory, getFunctionByName } from "../utils.js";
-import type { OpenAIToolCall } from "./types/chat.js";
 
 async function callFunction(
   tools: BaseTool[],
-  toolCall: OpenAIToolCall,
-): Promise<[ChatMessage, ToolOutput]> {
-  const id_ = toolCall.id;
-  const functionCall = toolCall.function;
-  const name = toolCall.function.name;
-  const argumentsStr = toolCall.function.arguments;
+  toolCall: ToolCallOptions["toolCall"],
+): Promise<[ChatMessage<ToolCallLLMMessageOptions>, ToolOutput]> {
+  const id = toolCall.id;
+  const name = toolCall.name;
+  const input = toolCall.input;
 
   if (Settings.debug) {
     console.log("=== Calling Function ===");
-    console.log(`Calling function: ${name} with args: ${argumentsStr}`);
+    console.log(`Calling function: ${name} with args: ${input}`);
   }
 
   const tool = getFunctionByName(tools, name);
-  const argumentDict = JSON.parse(argumentsStr);
 
   // Call tool
   // Use default error message
-  const output = await callToolWithErrorHandling(tool, argumentDict, null);
+  const output = await callToolWithErrorHandling(tool, input);
 
   if (Settings.debug) {
     console.log(`Got output ${output}`);
@@ -58,11 +53,13 @@ async function callFunction(
 
   return [
     {
-      content: String(output),
-      role: "tool",
-      additionalKwargs: {
-        name,
-        tool_call_id: id_,
+      content: `${output}`,
+      role: "user",
+      options: {
+        toolResult: {
+          id,
+          isError: false,
+        },
       },
     },
     output,
@@ -74,7 +71,7 @@ type OpenAIAgentWorkerParams = {
   llm?: OpenAI;
   prefixMessages?: ChatMessage[];
   maxFunctionCalls?: number;
-  toolRetriever?: ObjectRetriever;
+  toolRetriever?: ObjectRetriever<BaseTool>;
 };
 
 type CallFunctionOutput = {
@@ -110,9 +107,9 @@ export class OpenAIAgentWorker
     }
     this.prefixMessages = prefixMessages || [];
 
-    if (Array.isArray(tools) && tools.length > 0 && toolRetriever) {
+    if (tools.length > 0 && toolRetriever) {
       throw new Error("Cannot specify both tools and tool_retriever");
-    } else if (Array.isArray(tools)) {
+    } else if (tools.length > 0) {
       this._getTools = async () => tools;
     } else if (toolRetriever) {
       // fixme: this won't work, type mismatch
@@ -123,7 +120,7 @@ export class OpenAIAgentWorker
     }
   }
 
-  public getAllMessages(task: Task): ChatMessage[] {
+  public getAllMessages(task: Task): ChatMessage<ToolCallLLMMessageOptions>[] {
     return [
       ...this.prefixMessages,
       ...task.memory.get(),
@@ -131,29 +128,33 @@ export class OpenAIAgentWorker
     ];
   }
 
-  public getLatestToolCalls(task: Task): OpenAIToolCall[] | null {
+  public getLatestToolCall(task: Task): ToolCallOptions["toolCall"] | null {
     const chatHistory: ChatMessage[] = task.extraState.newMemory.getAll();
 
     if (chatHistory.length === 0) {
       return null;
     }
 
-    return chatHistory[chatHistory.length - 1].additionalKwargs?.toolCalls;
+    // @ts-expect-error fixme
+    return chatHistory[chatHistory.length - 1].options?.toolCall;
   }
 
   private _getLlmChatParams(
     task: Task,
-    openaiTools: BaseTool[],
+    tools: BaseTool[],
     toolChoice: ChatCompletionToolChoiceOption = "auto",
-  ): LLMChatParamsBase<OpenAIAdditionalChatOptions> {
+  ): LLMChatParamsBase<OpenAIAdditionalChatOptions, ToolCallLLMMessageOptions> {
     const llmChatParams = {
       messages: this.getAllMessages(task),
-      tools: [] as BaseTool[],
+      tools: undefined as BaseTool[] | undefined,
       additionalChatOptions: {} as OpenAIAdditionalChatOptions,
-    } satisfies LLMChatParamsBase<OpenAIAdditionalChatOptions>;
+    } satisfies LLMChatParamsBase<
+      OpenAIAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >;
 
-    if (openaiTools.length > 0) {
-      llmChatParams.tools = openaiTools;
+    if (tools.length > 0) {
+      llmChatParams.tools = tools;
       llmChatParams.additionalChatOptions.tool_choice = toolChoice;
     }
 
@@ -174,52 +175,88 @@ export class OpenAIAgentWorker
 
   private async _getStreamAiResponse(
     task: Task,
-    llmChatParams: LLMChatParamsBase<OpenAIAdditionalChatOptions>,
+    llmChatParams: LLMChatParamsBase<
+      OpenAIAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >,
   ): Promise<StreamingAgentChatResponse | AgentChatResponse> {
     const stream = await this.llm.chat({
       stream: true,
       ...llmChatParams,
     });
-    // read first chunk from stream to find out if we need to call tools
-    const iterator = stream[Symbol.asyncIterator]();
-    let { value } = await iterator.next();
-    let content = value.delta;
-    const hasToolCalls = value.additionalKwargs?.toolCalls.length > 0;
 
-    if (hasToolCalls) {
-      // consume stream until we have all the tool calls and return a non-streamed response
-      for await (value of stream) {
-        content += value.delta;
-      }
-      return this._processMessage(task, {
-        content,
-        role: "assistant",
-        additionalKwargs: value.additionalKwargs,
-      });
+    const responseChunkStream = new ReadableStream<
+      ChatResponseChunk<ToolCallLLMMessageOptions>
+    >({
+      async start(controller) {
+        for await (const chunk of stream) {
+          controller.enqueue(chunk);
+        }
+        controller.close();
+      },
+    });
+    const [pipStream, finalStream] = responseChunkStream.tee();
+    const reader = pipStream.getReader();
+    const { value } = await reader.read();
+    reader.releaseLock();
+    if (value === undefined) {
+      throw new Error("first chunk value is undefined, this should not happen");
     }
-
-    const newStream = streamConverter.bind(this)(
-      streamReducer({
-        stream,
-        initialValue: content,
-        reducer: (accumulator, part) => (accumulator += part.delta),
-        finished: (accumulator) => {
-          task.extraState.newMemory.put({
-            content: accumulator,
-            role: "assistant",
-          });
-        },
-      }),
-      (r: ChatResponseChunk) => new Response(r.delta),
+    // check if first chunk has tool calls, if so, this is a function call
+    // otherwise, it's a regular message
+    const hasToolCall: boolean = !!(
+      value.options && "toolCall" in value.options
     );
 
-    return new StreamingAgentChatResponse(newStream, task.extraState.sources);
+    if (hasToolCall) {
+      return this._processMessage(task, {
+        content: await pipeline(finalStream, async (iterator) => {
+          let content = "";
+          for await (const value of iterator) {
+            content += value.delta;
+          }
+          return content;
+        }),
+        role: "assistant",
+        options: value.options,
+      });
+    } else {
+      const [responseStream, chunkStream] = finalStream.tee();
+      let content = "";
+      return new StreamingAgentChatResponse(
+        responseStream.pipeThrough<Response>({
+          readable: new ReadableStream({
+            async start(controller) {
+              for await (const chunk of chunkStream) {
+                controller.enqueue(new Response(chunk.delta));
+              }
+              controller.close();
+            },
+          }),
+          writable: new WritableStream({
+            write(chunk) {
+              content += chunk.delta;
+            },
+            close() {
+              task.extraState.newMemory.put({
+                content,
+                role: "assistant",
+              });
+            },
+          }),
+        }),
+        task.extraState.sources,
+      );
+    }
   }
 
   private async _getAgentResponse(
     task: Task,
     mode: ChatResponseMode,
-    llmChatParams: LLMChatParamsBase<OpenAIAdditionalChatOptions>,
+    llmChatParams: LLMChatParamsBase<
+      OpenAIAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >,
   ): Promise<AgentChatResponse | StreamingAgentChatResponse> {
     if (mode === ChatResponseMode.WAIT) {
       const chatResponse = await this.llm.chat({
@@ -240,14 +277,8 @@ export class OpenAIAgentWorker
 
   async callFunction(
     tools: BaseTool[],
-    toolCall: OpenAIToolCall,
+    toolCall: ToolCallOptions["toolCall"],
   ): Promise<CallFunctionOutput> {
-    const functionCall = toolCall.function;
-
-    if (!functionCall) {
-      throw new Error("Invalid tool_call object");
-    }
-
     const functionMessage = await callFunction(tools, toolCall);
 
     const message = functionMessage[0];
@@ -281,18 +312,14 @@ export class OpenAIAgentWorker
   }
 
   private _shouldContinue(
-    toolCalls: OpenAIToolCall[] | null,
+    toolCall: ToolCallOptions["toolCall"] | null,
     nFunctionCalls: number,
-  ): boolean {
+  ): toolCall is ToolCallOptions["toolCall"] {
     if (nFunctionCalls > this.maxFunctionCalls) {
       return false;
     }
 
-    if (toolCalls?.length === 0) {
-      return false;
-    }
-
-    return true;
+    return !!toolCall;
   }
 
   async getTools(input: string): Promise<BaseTool[]> {
@@ -319,29 +346,25 @@ export class OpenAIAgentWorker
       llmChatParams,
     );
 
-    const latestToolCalls = this.getLatestToolCalls(task) || [];
+    const latestToolCall = this.getLatestToolCall(task) ?? null;
 
     let isDone: boolean;
-    let newSteps: TaskStep[] = [];
+    let newSteps: TaskStep[];
 
-    if (
-      !this._shouldContinue(latestToolCalls, task.extraState.nFunctionCalls)
-    ) {
+    if (!this._shouldContinue(latestToolCall, task.extraState.nFunctionCalls)) {
       isDone = true;
       newSteps = [];
     } else {
       isDone = false;
-      for (const toolCall of latestToolCalls) {
-        const { message, toolOutput } = await this.callFunction(
-          tools,
-          toolCall,
-        );
+      const { message, toolOutput } = await this.callFunction(
+        tools,
+        latestToolCall,
+      );
 
-        task.extraState.sources.push(toolOutput);
-        task.extraState.newMemory.put(message);
+      task.extraState.sources.push(toolOutput);
+      task.extraState.newMemory.put(message);
 
-        task.extraState.nFunctionCalls += 1;
-      }
+      task.extraState.nFunctionCalls += 1;
 
       newSteps = [step.getNextStep(randomUUID(), undefined)];
     }

packages/core/src/agent/react/base.ts

@@ -1,4 +1,5 @@
 import type { ChatMessage, LLM } from "../../llm/index.js";
+import type { BaseMemory } from "../../memory/types.js";
 import type { ObjectRetriever } from "../../objects/base.js";
 import type { BaseTool } from "../../types.js";
 import { AgentRunner } from "../runner/base.js";
@@ -7,11 +8,11 @@ import { ReActAgentWorker } from "./worker.js";
 type ReActAgentParams = {
   tools: BaseTool[];
   llm?: LLM;
-  memory?: any;
+  memory?: BaseMemory;
   prefixMessages?: ChatMessage[];
   maxInteractions?: number;
   defaultToolChoice?: string;
-  toolRetriever?: ObjectRetriever;
+  toolRetriever?: ObjectRetriever<BaseTool>;
 };
 
 /**
@@ -40,6 +41,7 @@ export class ReActAgent extends AgentRunner {
       agentWorker: stepEngine,
       memory,
       defaultToolChoice,
+      // @ts-expect-error 2322
       chatHistory: prefixMessages,
     });
   }

packages/core/src/agent/react/worker.ts

@@ -2,6 +2,7 @@ import { randomUUID } from "@llamaindex/env";
 import type { ChatMessage } from "cohere-ai/api";
 import { Settings } from "../../Settings.js";
 import { AgentChatResponse } from "../../engines/chat/index.js";
+import { getCallbackManager } from "../../internal/settings/CallbackManager.js";
 import { type ChatResponse, type LLM } from "../../llm/index.js";
 import { extractText } from "../../llm/utils.js";
 import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer.js";
@@ -25,7 +26,7 @@ type ReActAgentWorkerParams = {
   maxInteractions?: number;
   reactChatFormatter?: ReActChatFormatter | undefined;
   outputParser?: ReActOutputParser | undefined;
-  toolRetriever?: ObjectRetriever | undefined;
+  toolRetriever?: ObjectRetriever<BaseTool> | undefined;
 };
 
 function addUserStepToReasoning(
@@ -194,7 +195,15 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
 
     const tool = toolsDict[actionReasoningStep.action];
 
-    const toolOutput = await tool?.call?.(actionReasoningStep.actionInput);
+    getCallbackManager().dispatchEvent("llm-tool-call", {
+      payload: {
+        toolCall: {
+          name: tool.metadata.name,
+          input: JSON.stringify(actionReasoningStep.actionInput),
+        },
+      },
+    });
+    const toolOutput = await tool.call!(actionReasoningStep.actionInput);
 
     task.extraState.sources.push(
       new ToolOutput(

packages/core/src/agent/runner/base.ts

@@ -1,11 +1,12 @@
 import { randomUUID } from "@llamaindex/env";
+import type { ChatHistory } from "../../ChatHistory.js";
 import type { ChatEngineAgentParams } from "../../engines/chat/index.js";
 import {
   AgentChatResponse,
   ChatResponseMode,
   StreamingAgentChatResponse,
 } from "../../engines/chat/index.js";
-import type { ChatMessage, LLM } from "../../llm/index.js";
+import type { LLM } from "../../llm/index.js";
 import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer.js";
 import type { BaseMemory } from "../../memory/types.js";
 import type { AgentWorker, TaskStepOutput } from "../types.js";
@@ -30,7 +31,7 @@ const validateStepFromArgs = (
 
 type AgentRunnerParams = {
   agentWorker: AgentWorker;
-  chatHistory?: ChatMessage[];
+  chatHistory?: ChatHistory;
   state?: AgentState;
   memory?: BaseMemory;
   llm?: LLM;
@@ -161,7 +162,7 @@ export class AgentRunner extends BaseAgentRunner {
     const task = this.state.getTask(taskId);
     const curStep = step || this.state.getStepQueue(taskId).shift();
 
-    let curStepOutput;
+    let curStepOutput: TaskStepOutput;
 
     if (!curStep) {
       throw new Error(`No step found for task ${taskId}`);

packages/core/src/agent/types.ts

@@ -4,6 +4,7 @@ import type {
   StreamingAgentChatResponse,
 } from "../engines/chat/index.js";
 
+import type { BaseMemory } from "../memory/types.js";
 import type { QueryEngineParamsNonStreaming } from "../types.js";
 
 export interface AgentWorker<ExtraParams extends object = object> {
@@ -72,7 +73,7 @@ export abstract class BaseAgent implements BaseChatEngine, BaseQueryEngine {
 type TaskParams = {
   taskId: string;
   input: string;
-  memory: any;
+  memory: BaseMemory;
   extraState: Record<string, any>;
 };
 
@@ -84,7 +85,7 @@ export class Task {
   taskId: string;
   input: string;
 
-  memory: any;
+  memory: BaseMemory;
   extraState: Record<string, any>;
 
   constructor({ taskId, input, memory, extraState }: TaskParams) {

packages/core/src/callbacks/CallbackManager.ts

@@ -5,7 +5,12 @@ import {
   EventCaller,
   getEventCaller,
 } from "../internal/context/EventCaller.js";
-import type { LLMEndEvent, LLMStartEvent } from "../llm/types.js";
+import type {
+  LLMEndEvent,
+  LLMStartEvent,
+  LLMStreamEvent,
+  LLMToolCallEvent,
+} from "../llm/types.js";
 
 export class LlamaIndexCustomEvent<T = any> extends CustomEvent<T> {
   reason: EventCaller | null;
@@ -44,6 +49,8 @@ export interface LlamaIndexEventMaps {
   stream: CustomEvent<StreamCallbackResponse>;
   "llm-start": LLMStartEvent;
   "llm-end": LLMEndEvent;
+  "llm-tool-call": LLMToolCallEvent;
+  "llm-stream": LLMStreamEvent;
 }
 
 //#region @deprecated remove in the next major version
@@ -198,8 +205,10 @@ export class CallbackManager implements CallbackManagerMethods {
     if (!handlers) {
       return;
     }
-    handlers.forEach((handler) =>
-      handler(LlamaIndexCustomEvent.fromEvent(event, detail)),
-    );
+    queueMicrotask(() => {
+      handlers.forEach((handler) =>
+        handler(LlamaIndexCustomEvent.fromEvent(event, detail)),
+      );
+    });
   }
 }

packages/core/src/embeddings/MultiModalEmbedding.ts

@@ -1,5 +1,10 @@
-import type { ImageType } from "../Node.js";
-import { BaseEmbedding } from "./types.js";
+import {
+  MetadataMode,
+  splitNodesByType,
+  type BaseNode,
+  type ImageType,
+} from "../Node.js";
+import { BaseEmbedding, batchEmbeddings } from "./types.js";
 
 /*
  * Base class for Multi Modal embeddings.
@@ -8,9 +13,39 @@ import { BaseEmbedding } from "./types.js";
 export abstract class MultiModalEmbedding extends BaseEmbedding {
   abstract getImageEmbedding(images: ImageType): Promise<number[]>;
 
+  /**
+   * Optionally override this method to retrieve multiple image embeddings in a single request
+   * @param texts
+   */
   async getImageEmbeddings(images: ImageType[]): Promise<number[][]> {
     return Promise.all(
       images.map((imgFilePath) => this.getImageEmbedding(imgFilePath)),
     );
   }
+
+  async transform(nodes: BaseNode[], _options?: any): Promise<BaseNode[]> {
+    const { imageNodes, textNodes } = splitNodesByType(nodes);
+
+    const embeddings = await batchEmbeddings(
+      textNodes.map((node) => node.getContent(MetadataMode.EMBED)),
+      this.getTextEmbeddings.bind(this),
+      this.embedBatchSize,
+      _options,
+    );
+    for (let i = 0; i < textNodes.length; i++) {
+      textNodes[i].embedding = embeddings[i];
+    }
+
+    const imageEmbeddings = await batchEmbeddings(
+      imageNodes.map((n) => n.image),
+      this.getImageEmbeddings.bind(this),
+      this.embedBatchSize,
+      _options,
+    );
+    for (let i = 0; i < imageNodes.length; i++) {
+      imageNodes[i].embedding = imageEmbeddings[i];
+    }
+
+    return nodes;
+  }
 }

packages/core/src/embeddings/types.ts

@@ -5,6 +5,8 @@ import { SimilarityType, similarity } from "./utils.js";
 
 const DEFAULT_EMBED_BATCH_SIZE = 10;
 
+type EmbedFunc<T> = (values: T[]) => Promise<Array<number[]>>;
+
 export abstract class BaseEmbedding implements TransformComponent {
   embedBatchSize = DEFAULT_EMBED_BATCH_SIZE;
 
@@ -45,35 +47,18 @@ export abstract class BaseEmbedding implements TransformComponent {
       logProgress?: boolean;
     },
   ): Promise<Array<number[]>> {
-    const resultEmbeddings: Array<number[]> = [];
-    const chunkSize = this.embedBatchSize;
-
-    const queue: string[] = texts;
-
-    const curBatch: string[] = [];
-
-    for (let i = 0; i < queue.length; i++) {
-      curBatch.push(queue[i]);
-      if (i == queue.length - 1 || curBatch.length == chunkSize) {
-        const embeddings = await this.getTextEmbeddings(curBatch);
-
-        resultEmbeddings.push(...embeddings);
-
-        if (options?.logProgress) {
-          console.log(`getting embedding progress: ${i} / ${queue.length}`);
-        }
-
-        curBatch.length = 0;
-      }
-    }
-
-    return resultEmbeddings;
+    return await batchEmbeddings(
+      texts,
+      this.getTextEmbeddings.bind(this),
+      this.embedBatchSize,
+      options,
+    );
   }
 
   async transform(nodes: BaseNode[], _options?: any): Promise<BaseNode[]> {
     const texts = nodes.map((node) => node.getContent(MetadataMode.EMBED));
 
-    const embeddings = await this.getTextEmbeddingsBatch(texts);
+    const embeddings = await this.getTextEmbeddingsBatch(texts, _options);
 
     for (let i = 0; i < nodes.length; i++) {
       nodes[i].embedding = embeddings[i];
@@ -82,3 +67,35 @@ export abstract class BaseEmbedding implements TransformComponent {
     return nodes;
   }
 }
+
+export async function batchEmbeddings<T>(
+  values: T[],
+  embedFunc: EmbedFunc<T>,
+  chunkSize: number,
+  options?: {
+    logProgress?: boolean;
+  },
+): Promise<Array<number[]>> {
+  const resultEmbeddings: Array<number[]> = [];
+
+  const queue: T[] = values;
+
+  const curBatch: T[] = [];
+
+  for (let i = 0; i < queue.length; i++) {
+    curBatch.push(queue[i]);
+    if (i == queue.length - 1 || curBatch.length == chunkSize) {
+      const embeddings = await embedFunc(curBatch);
+
+      resultEmbeddings.push(...embeddings);
+
+      if (options?.logProgress) {
+        console.log(`getting embedding progress: ${i} / ${queue.length}`);
+      }
+
+      curBatch.length = 0;
+    }
+  }
+
+  return resultEmbeddings;
+}

packages/core/src/engines/chat/ContextChatEngine.ts

@@ -113,10 +113,9 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
     });
     const textOnly = extractText(message);
     const context = await this.contextGenerator.generate(textOnly);
-    const nodes = context.nodes.map((r) => r.node);
     const messages = await chatHistory.requestMessages(
       context ? [context.message] : undefined,
     );
-    return { nodes, messages };
+    return { nodes: context.nodes, messages };
   }
 }

packages/core/src/engines/chat/types.ts

@@ -1,5 +1,5 @@
 import type { ChatHistory } from "../../ChatHistory.js";
-import type { BaseNode, NodeWithScore } from "../../Node.js";
+import type { NodeWithScore } from "../../Node.js";
 import type { Response } from "../../Response.js";
 import type { ChatMessage } from "../../llm/index.js";
 import type { MessageContent } from "../../llm/types.js";
@@ -66,12 +66,12 @@ export enum ChatResponseMode {
 export class AgentChatResponse {
   response: string;
   sources: ToolOutput[];
-  sourceNodes?: BaseNode[];
+  sourceNodes?: NodeWithScore[];
 
   constructor(
     response: string,
     sources?: ToolOutput[],
-    sourceNodes?: BaseNode[],
+    sourceNodes?: NodeWithScore[],
   ) {
     this.response = response;
     this.sources = sources || [];
@@ -91,12 +91,12 @@ export class StreamingAgentChatResponse {
   response: AsyncIterable<Response>;
 
   sources: ToolOutput[];
-  sourceNodes?: BaseNode[];
+  sourceNodes?: NodeWithScore[];
 
   constructor(
     response: AsyncIterable<Response>,
     sources?: ToolOutput[],
-    sourceNodes?: BaseNode[],
+    sourceNodes?: NodeWithScore[],
   ) {
     this.response = response;
     this.sources = sources ?? [];

packages/core/src/engines/query/RouterQueryEngine.ts

@@ -1,4 +1,4 @@
-import type { BaseNode } from "../../Node.js";
+import type { NodeWithScore } from "../../Node.js";
 import { Response } from "../../Response.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { llmFromSettingsOrContext } from "../../Settings.js";
@@ -33,7 +33,7 @@ async function combineResponses(
   }
 
   const responseStrs: string[] = [];
-  const sourceNodes: BaseNode[] = [];
+  const sourceNodes: NodeWithScore[] = [];
 
   for (const response of responses) {
     if (response?.sourceNodes) {

packages/core/src/evaluation/Correctness.ts

@@ -111,7 +111,7 @@ export class CorrectnessEvaluator extends PromptMixin implements BaseEvaluator {
 
     if (response) {
       for (const node of response.sourceNodes || []) {
-        contexts.push(node.getContent(MetadataMode.ALL));
+        contexts.push(node.node.getContent(MetadataMode.ALL));
       }
     }
 

packages/core/src/evaluation/Faithfulness.ts

@@ -137,7 +137,7 @@ export class FaithfulnessEvaluator
 
     if (response) {
       for (const node of response.sourceNodes || []) {
-        contexts.push(node.getContent(MetadataMode.ALL));
+        contexts.push(node.node.getContent(MetadataMode.ALL));
       }
     }
 

packages/core/src/evaluation/Relevancy.ts

@@ -126,7 +126,7 @@ export class RelevancyEvaluator extends PromptMixin implements BaseEvaluator {
 
     if (response) {
       for (const node of response.sourceNodes || []) {
-        contexts.push(node.getContent(MetadataMode.ALL));
+        contexts.push(node.node.getContent(MetadataMode.ALL));
       }
     }
 

packages/core/src/indices/vectorStore/index.ts

@@ -4,12 +4,7 @@ import type {
   Metadata,
   NodeWithScore,
 } from "../../Node.js";
-import {
-  ImageNode,
-  MetadataMode,
-  ObjectType,
-  splitNodesByType,
-} from "../../Node.js";
+import { ImageNode, ObjectType, splitNodesByType } from "../../Node.js";
 import type { BaseRetriever, RetrieveParams } from "../../Retriever.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import {
@@ -179,14 +174,21 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
     nodes: BaseNode[],
     options?: { logProgress?: boolean },
   ): Promise<BaseNode[]> {
-    const texts = nodes.map((node) => node.getContent(MetadataMode.EMBED));
-    const embeddings = await this.embedModel.getTextEmbeddingsBatch(texts, {
+    const { imageNodes, textNodes } = splitNodesByType(nodes);
+    if (imageNodes.length > 0) {
+      if (!this.imageEmbedModel) {
+        throw new Error(
+          "Cannot calculate image nodes embedding without 'imageEmbedModel' set",
+        );
+      }
+      await this.imageEmbedModel.transform(imageNodes, {
+        logProgress: options?.logProgress,
+      });
+    }
+    await this.embedModel.transform(textNodes, {
       logProgress: options?.logProgress,
     });
-    return nodes.map((node, i) => {
-      node.embedding = embeddings[i];
-      return node;
-    });
+    return nodes;
   }
 
   /**
@@ -324,25 +326,15 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
     if (!nodes || nodes.length === 0) {
       return;
     }
+    nodes = await this.getNodeEmbeddingResults(nodes, options);
     const { imageNodes, textNodes } = splitNodesByType(nodes);
     if (imageNodes.length > 0) {
       if (!this.imageVectorStore) {
         throw new Error("Cannot insert image nodes without image vector store");
       }
-      const imageNodesWithEmbedding = await this.getImageNodeEmbeddingResults(
-        imageNodes,
-        options,
-      );
-      await this.insertNodesToStore(
-        this.imageVectorStore,
-        imageNodesWithEmbedding,
-      );
+      await this.insertNodesToStore(this.imageVectorStore, imageNodes);
     }
-    const embeddingResults = await this.getNodeEmbeddingResults(
-      textNodes,
-      options,
-    );
-    await this.insertNodesToStore(this.vectorStore, embeddingResults);
+    await this.insertNodesToStore(this.vectorStore, textNodes);
     await this.indexStore.addIndexStruct(this.indexStruct);
   }
 
@@ -378,35 +370,6 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
       await this.indexStore.addIndexStruct(this.indexStruct);
     }
   }
-
-  /**
-   * Calculates the embeddings for the given image nodes.
-   *
-   * @param nodes - An array of ImageNode objects representing the nodes for which embeddings are to be calculated.
-   * @param {Object} [options] - An optional object containing additional parameters.
-   *   @param {boolean} [options.logProgress] - A boolean indicating whether to log progress to the console (useful for debugging).
-   */
-  async getImageNodeEmbeddingResults(
-    nodes: ImageNode[],
-    options?: { logProgress?: boolean },
-  ): Promise<ImageNode[]> {
-    if (!this.imageEmbedModel) {
-      return [];
-    }
-
-    const nodesWithEmbeddings: ImageNode[] = [];
-
-    for (let i = 0; i < nodes.length; ++i) {
-      const node = nodes[i];
-      if (options?.logProgress) {
-        console.log(`Getting embedding for node ${i + 1}/${nodes.length}`);
-      }
-      node.embedding = await this.imageEmbedModel.getImageEmbedding(node.image);
-      nodesWithEmbeddings.push(node);
-    }
-
-    return nodesWithEmbeddings;
-  }
 }
 
 /**

packages/core/src/internal/utils.ts

@@ -5,3 +5,14 @@ export const isAsyncGenerator = (obj: unknown): obj is AsyncGenerator => {
 export const isGenerator = (obj: unknown): obj is Generator => {
   return obj != null && typeof obj === "object" && Symbol.iterator in obj;
 };
+
+/**
+ * Prettify an error for AI to read
+ */
+export function prettifyError(error: unknown): string {
+  if (error instanceof Error) {
+    return `Error(${error.name}): ${error.message}`;
+  } else {
+    return `${error}`;
+  }
+}

packages/core/src/llm/LLM.ts

@@ -1,372 +0,0 @@
-import { type StreamCallbackResponse } from "../callbacks/CallbackManager.js";
-
-import type { LLMOptions } from "portkey-ai";
-import { getCallbackManager } from "../internal/settings/CallbackManager.js";
-import { BaseLLM } from "./base.js";
-import type { PortkeySession } from "./portkey.js";
-import { getPortkeySession } from "./portkey.js";
-import { ReplicateSession } from "./replicate_ai.js";
-import type {
-  ChatMessage,
-  ChatResponse,
-  ChatResponseChunk,
-  LLMChatParamsNonStreaming,
-  LLMChatParamsStreaming,
-  LLMMetadata,
-  MessageType,
-} from "./types.js";
-import { extractText, wrapLLMEvent } from "./utils.js";
-
-export const ALL_AVAILABLE_LLAMADEUCE_MODELS = {
-  "Llama-2-70b-chat-old": {
-    contextWindow: 4096,
-    replicateApi:
-      "replicate/llama70b-v2-chat:e951f18578850b652510200860fc4ea62b3b16fac280f83ff32282f87bbd2e48",
-    //^ Previous 70b model. This is also actually 4 bit, although not exllama.
-  },
-  "Llama-2-70b-chat-4bit": {
-    contextWindow: 4096,
-    replicateApi:
-      "meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3",
-    //^ Model is based off of exllama 4bit.
-  },
-  "Llama-2-13b-chat-old": {
-    contextWindow: 4096,
-    replicateApi:
-      "a16z-infra/llama13b-v2-chat:df7690f1994d94e96ad9d568eac121aecf50684a0b0963b25a41cc40061269e5",
-  },
-  //^ Last known good 13b non-quantized model. In future versions they add the SYS and INST tags themselves
-  "Llama-2-13b-chat-4bit": {
-    contextWindow: 4096,
-    replicateApi:
-      "meta/llama-2-13b-chat:f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
-  },
-  "Llama-2-7b-chat-old": {
-    contextWindow: 4096,
-    replicateApi:
-      "a16z-infra/llama7b-v2-chat:4f0a4744c7295c024a1de15e1a63c880d3da035fa1f49bfd344fe076074c8eea",
-    //^ Last (somewhat) known good 7b non-quantized model. In future versions they add the SYS and INST
-    // tags themselves
-    // https://github.com/replicate/cog-llama-template/commit/fa5ce83912cf82fc2b9c01a4e9dc9bff6f2ef137
-    // Problem is that they fix the max_new_tokens issue in the same commit. :-(
-  },
-  "Llama-2-7b-chat-4bit": {
-    contextWindow: 4096,
-    replicateApi:
-      "meta/llama-2-7b-chat:13c3cdee13ee059ab779f0291d29054dab00a47dad8261375654de5540165fb0",
-  },
-};
-
-export enum DeuceChatStrategy {
-  A16Z = "a16z",
-  META = "meta",
-  METAWBOS = "metawbos",
-  //^ This is not exactly right because SentencePiece puts the BOS and EOS token IDs in after tokenization
-  // Unfortunately any string only API won't support these properly.
-  REPLICATE4BIT = "replicate4bit",
-  //^ To satisfy Replicate's 4 bit models' requirements where they also insert some INST tags
-  REPLICATE4BITWNEWLINES = "replicate4bitwnewlines",
-  //^ Replicate's documentation recommends using newlines: https://replicate.com/blog/how-to-prompt-llama
-}
-
-/**
- * Llama2 LLM implementation
- */
-export class LlamaDeuce extends BaseLLM {
-  model: keyof typeof ALL_AVAILABLE_LLAMADEUCE_MODELS;
-  chatStrategy: DeuceChatStrategy;
-  temperature: number;
-  topP: number;
-  maxTokens?: number;
-  replicateSession: ReplicateSession;
-
-  constructor(init?: Partial<LlamaDeuce>) {
-    super();
-    this.model = init?.model ?? "Llama-2-70b-chat-4bit";
-    this.chatStrategy =
-      init?.chatStrategy ??
-      (this.model.endsWith("4bit")
-        ? DeuceChatStrategy.REPLICATE4BITWNEWLINES // With the newer Replicate models they do the system message themselves.
-        : DeuceChatStrategy.METAWBOS); // With BOS and EOS seems to work best, although they all have problems past a certain point
-    this.temperature = init?.temperature ?? 0.1; // minimum temperature is 0.01 for Replicate endpoint
-    this.topP = init?.topP ?? 1;
-    this.maxTokens =
-      init?.maxTokens ??
-      ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model].contextWindow; // For Replicate, the default is 500 tokens which is too low.
-    this.replicateSession = init?.replicateSession ?? new ReplicateSession();
-  }
-
-  get metadata() {
-    return {
-      model: this.model,
-      temperature: this.temperature,
-      topP: this.topP,
-      maxTokens: this.maxTokens,
-      contextWindow: ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model].contextWindow,
-      tokenizer: undefined,
-    };
-  }
-
-  mapMessagesToPrompt(messages: ChatMessage[]) {
-    if (this.chatStrategy === DeuceChatStrategy.A16Z) {
-      return this.mapMessagesToPromptA16Z(messages);
-    } else if (this.chatStrategy === DeuceChatStrategy.META) {
-      return this.mapMessagesToPromptMeta(messages);
-    } else if (this.chatStrategy === DeuceChatStrategy.METAWBOS) {
-      return this.mapMessagesToPromptMeta(messages, { withBos: true });
-    } else if (this.chatStrategy === DeuceChatStrategy.REPLICATE4BIT) {
-      return this.mapMessagesToPromptMeta(messages, {
-        replicate4Bit: true,
-        withNewlines: true,
-      });
-    } else if (this.chatStrategy === DeuceChatStrategy.REPLICATE4BITWNEWLINES) {
-      return this.mapMessagesToPromptMeta(messages, {
-        replicate4Bit: true,
-        withNewlines: true,
-      });
-    } else {
-      return this.mapMessagesToPromptMeta(messages);
-    }
-  }
-
-  mapMessagesToPromptA16Z(messages: ChatMessage[]) {
-    return {
-      prompt:
-        messages.reduce((acc, message) => {
-          return (
-            (acc && `${acc}\n\n`) +
-            `${this.mapMessageTypeA16Z(message.role)}${message.content}`
-          );
-        }, "") + "\n\nAssistant:",
-      //^ Here we're differing from A16Z by omitting the space. Generally spaces at the end of prompts decrease performance due to tokenization
-      systemPrompt: undefined,
-    };
-  }
-
-  mapMessageTypeA16Z(messageType: MessageType): string {
-    switch (messageType) {
-      case "user":
-        return "User: ";
-      case "assistant":
-        return "Assistant: ";
-      case "system":
-        return "";
-      default:
-        throw new Error("Unsupported LlamaDeuce message type");
-    }
-  }
-
-  mapMessagesToPromptMeta(
-    messages: ChatMessage[],
-    opts?: {
-      withBos?: boolean;
-      replicate4Bit?: boolean;
-      withNewlines?: boolean;
-    },
-  ) {
-    const {
-      withBos = false,
-      replicate4Bit = false,
-      withNewlines = false,
-    } = opts ?? {};
-    const DEFAULT_SYSTEM_PROMPT = `You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
-
-If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.`;
-
-    const B_SYS = "<<SYS>>\n";
-    const E_SYS = "\n<</SYS>>\n\n";
-    const B_INST = "[INST]";
-    const E_INST = "[/INST]";
-    const BOS = "<s>";
-    const EOS = "</s>";
-
-    if (messages.length === 0) {
-      return { prompt: "", systemPrompt: undefined };
-    }
-
-    messages = [...messages]; // so we can use shift without mutating the original array
-
-    let systemPrompt = undefined;
-    if (messages[0].role === "system") {
-      const systemMessage = messages.shift()!;
-
-      if (replicate4Bit) {
-        systemPrompt = systemMessage.content;
-      } else {
-        const systemStr = `${B_SYS}${systemMessage.content}${E_SYS}`;
-
-        // TS Bug: https://github.com/microsoft/TypeScript/issues/9998
-        // @ts-ignore
-        if (messages[0].role !== "user") {
-          throw new Error(
-            "LlamaDeuce: if there is a system message, the second message must be a user message.",
-          );
-        }
-
-        const userContent = messages[0].content;
-
-        messages[0].content = `${systemStr}${userContent}`;
-      }
-    } else {
-      if (!replicate4Bit) {
-        messages[0].content = `${B_SYS}${DEFAULT_SYSTEM_PROMPT}${E_SYS}${messages[0].content}`;
-      }
-    }
-
-    return {
-      prompt: messages.reduce((acc, message, index) => {
-        const content = extractText(message.content);
-        if (index % 2 === 0) {
-          return (
-            `${acc}${withBos ? BOS : ""}${B_INST} ${content.trim()} ${E_INST}` +
-            (withNewlines ? "\n" : "")
-          );
-        } else {
-          return (
-            `${acc} ${content.trim()}` +
-            (withNewlines ? "\n" : " ") +
-            (withBos ? EOS : "")
-          ); // Yes, the EOS comes after the space. This is not a mistake.
-        }
-      }, ""),
-      systemPrompt,
-    };
-  }
-
-  chat(
-    params: LLMChatParamsStreaming,
-  ): Promise<AsyncIterable<ChatResponseChunk>>;
-  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
-  @wrapLLMEvent
-  async chat(
-    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
-  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, stream } = params;
-    const api = ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model]
-      .replicateApi as `${string}/${string}:${string}`;
-
-    const { prompt, systemPrompt } = this.mapMessagesToPrompt(messages);
-
-    const replicateOptions: any = {
-      input: {
-        prompt,
-        system_prompt: systemPrompt,
-        temperature: this.temperature,
-        top_p: this.topP,
-      },
-    };
-
-    if (this.model.endsWith("4bit")) {
-      replicateOptions.input.max_new_tokens = this.maxTokens;
-    } else {
-      replicateOptions.input.max_length = this.maxTokens;
-    }
-
-    //TODO: Add streaming for this
-    if (stream) {
-      throw new Error("Streaming not supported for LlamaDeuce");
-    }
-
-    //Non-streaming
-    const response = await this.replicateSession.replicate.run(
-      api,
-      replicateOptions,
-    );
-    return {
-      message: {
-        content: (response as Array<string>).join("").trimStart(),
-        //^ We need to do this because Replicate returns a list of strings (for streaming functionality which is not exposed by the run function)
-        role: "assistant",
-      },
-    };
-  }
-}
-
-export class Portkey extends BaseLLM {
-  apiKey?: string = undefined;
-  baseURL?: string = undefined;
-  mode?: string = undefined;
-  llms?: [LLMOptions] | null = undefined;
-  session: PortkeySession;
-
-  constructor(init?: Partial<Portkey>) {
-    super();
-    this.apiKey = init?.apiKey;
-    this.baseURL = init?.baseURL;
-    this.mode = init?.mode;
-    this.llms = init?.llms;
-    this.session = getPortkeySession({
-      apiKey: this.apiKey,
-      baseURL: this.baseURL,
-      llms: this.llms,
-      mode: this.mode,
-    });
-  }
-
-  get metadata(): LLMMetadata {
-    throw new Error("metadata not implemented for Portkey");
-  }
-
-  chat(
-    params: LLMChatParamsStreaming,
-  ): Promise<AsyncIterable<ChatResponseChunk>>;
-  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
-  @wrapLLMEvent
-  async chat(
-    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
-  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, stream, additionalChatOptions } = params;
-    if (stream) {
-      return this.streamChat(messages, additionalChatOptions);
-    } else {
-      const bodyParams = additionalChatOptions || {};
-      const response = await this.session.portkey.chatCompletions.create({
-        messages: messages.map((message) => ({
-          content: extractText(message.content),
-          role: message.role,
-        })),
-        ...bodyParams,
-      });
-
-      const content = response.choices[0].message?.content ?? "";
-      const role = response.choices[0].message?.role || "assistant";
-      return { message: { content, role: role as MessageType } };
-    }
-  }
-
-  async *streamChat(
-    messages: ChatMessage[],
-    params?: Record<string, any>,
-  ): AsyncIterable<ChatResponseChunk> {
-    const chunkStream = await this.session.portkey.chatCompletions.create({
-      messages: messages.map((message) => ({
-        content: extractText(message.content),
-        role: message.role,
-      })),
-      ...params,
-      stream: true,
-    });
-
-    //Indices
-    let idx_counter: number = 0;
-    for await (const part of chunkStream) {
-      //Increment
-      part.choices[0].index = idx_counter;
-      const is_done: boolean =
-        part.choices[0].finish_reason === "stop" ? true : false;
-      //onLLMStream Callback
-
-      const stream_callback: StreamCallbackResponse = {
-        index: idx_counter,
-        isDone: is_done,
-        // token: part,
-      };
-      getCallbackManager().dispatchEvent("stream", stream_callback);
-
-      idx_counter++;
-
-      yield { delta: part.choices[0].delta?.content ?? "" };
-    }
-    return;
-  }
-}

packages/core/src/llm/anthropic.ts

@@ -1,15 +1,30 @@
 import type { ClientOptions } from "@anthropic-ai/sdk";
 import { Anthropic as SDKAnthropic } from "@anthropic-ai/sdk";
+import type {
+  Tool,
+  ToolResultBlockParam,
+  ToolUseBlock,
+  ToolUseBlockParam,
+  ToolsBetaContentBlock,
+  ToolsBetaMessageParam,
+} from "@anthropic-ai/sdk/resources/beta/tools/messages";
+import type {
+  TextBlock,
+  TextBlockParam,
+} from "@anthropic-ai/sdk/resources/index";
+import type { MessageParam } from "@anthropic-ai/sdk/resources/messages";
 import { getEnv } from "@llamaindex/env";
+import _ from "lodash";
+import type { BaseTool } from "../types.js";
+import { ToolCallLLM } from "./base.js";
 import type {
   ChatMessage,
   ChatResponse,
   ChatResponseChunk,
   LLMChatParamsNonStreaming,
   LLMChatParamsStreaming,
-} from "llamaindex";
-import _ from "lodash";
-import { BaseLLM } from "./base.js";
+  ToolCallLLMMessageOptions,
+} from "./types.js";
 import { extractText, wrapLLMEvent } from "./utils.js";
 
 export class AnthropicSession {
@@ -81,7 +96,9 @@ const AVAILABLE_ANTHROPIC_MODELS_WITHOUT_DATE: { [key: string]: string } = {
   "claude-3-haiku": "claude-3-haiku-20240307",
 } as { [key in keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS]: string };
 
-export class Anthropic extends BaseLLM {
+export type AnthropicAdditionalChatOptions = {};
+
+export class Anthropic extends ToolCallLLM<AnthropicAdditionalChatOptions> {
   // Per completion Anthropic params
   model: keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS;
   temperature: number;
@@ -113,6 +130,10 @@ export class Anthropic extends BaseLLM {
       });
   }
 
+  get supportToolCall() {
+    return this.model.startsWith("claude-3");
+  }
+
   get metadata() {
     return {
       model: this.model,
@@ -131,30 +152,93 @@ export class Anthropic extends BaseLLM {
     return model;
   };
 
-  formatMessages(messages: ChatMessage[]) {
-    return messages.map((message) => {
+  formatMessages<Beta = false>(
+    messages: ChatMessage<ToolCallLLMMessageOptions>[],
+  ): Beta extends true ? ToolsBetaMessageParam[] : MessageParam[] {
+    return messages.map<any>((message) => {
       if (message.role !== "user" && message.role !== "assistant") {
         throw new Error("Unsupported Anthropic role");
       }
+      const options = message.options ?? {};
+      if ("toolResult" in options) {
+        const { id, isError } = options.toolResult;
+        return {
+          role: "user",
+          content: [
+            {
+              type: "tool_result",
+              is_error: isError,
+              content: [
+                {
+                  type: "text",
+                  text: extractText(message.content),
+                },
+              ],
+              tool_use_id: id,
+            },
+          ] satisfies ToolResultBlockParam[],
+        } satisfies ToolsBetaMessageParam;
+      } else if ("toolCall" in options) {
+        const aiThinkingText = extractText(message.content);
+        return {
+          role: "assistant",
+          content: [
+            // this could be empty when you call two tools in one query
+            ...(aiThinkingText.trim()
+              ? [
+                  {
+                    type: "text",
+                    text: aiThinkingText,
+                  } satisfies TextBlockParam,
+                ]
+              : []),
+            {
+              type: "tool_use",
+              id: options.toolCall.id,
+              name: options.toolCall.name,
+              input: options.toolCall.input,
+            } satisfies ToolUseBlockParam,
+          ] satisfies ToolsBetaContentBlock[],
+        } satisfies ToolsBetaMessageParam;
+      }
 
       return {
         content: extractText(message.content),
         role: message.role,
-      };
+      } satisfies MessageParam;
     });
   }
 
   chat(
-    params: LLMChatParamsStreaming,
-  ): Promise<AsyncIterable<ChatResponseChunk>>;
-  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+    params: LLMChatParamsStreaming<
+      AnthropicAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >,
+  ): Promise<AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>>>;
+  chat(
+    params: LLMChatParamsNonStreaming<
+      AnthropicAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >,
+  ): Promise<ChatResponse<ToolCallLLMMessageOptions>>;
   @wrapLLMEvent
   async chat(
-    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
-  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    params:
+      | LLMChatParamsNonStreaming<
+          AnthropicAdditionalChatOptions,
+          ToolCallLLMMessageOptions
+        >
+      | LLMChatParamsStreaming<
+          AnthropicAdditionalChatOptions,
+          ToolCallLLMMessageOptions
+        >,
+  ): Promise<
+    | ChatResponse<ToolCallLLMMessageOptions>
+    | AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>>
+  > {
     let { messages } = params;
 
-    const { stream } = params;
+    const { stream, tools } = params;
 
     let systemPrompt: string | null = null;
 
@@ -169,33 +253,80 @@ export class Anthropic extends BaseLLM {
       messages = messages.filter((message) => message.role !== "system");
     }
 
-    //Streaming
+    // case: Streaming
     if (stream) {
+      if (tools) {
+        console.error("Tools are not supported in streaming mode");
+      }
       return this.streamChat(messages, systemPrompt);
     }
+    // case: Non-streaming
+    const anthropic = this.session.anthropic;
 
-    //Non-streaming
-    const response = await this.session.anthropic.messages.create({
-      model: this.getModelName(this.model),
-      messages: this.formatMessages(messages),
-      max_tokens: this.maxTokens ?? 4096,
-      temperature: this.temperature,
-      top_p: this.topP,
-      ...(systemPrompt && { system: systemPrompt }),
-    });
+    if (tools) {
+      const response = await anthropic.beta.tools.messages.create({
+        messages: this.formatMessages<true>(messages),
+        tools: tools.map(Anthropic.toTool),
+        model: this.getModelName(this.model),
+        temperature: this.temperature,
+        max_tokens: this.maxTokens ?? 4096,
+        top_p: this.topP,
+        ...(systemPrompt && { system: systemPrompt }),
+      });
 
-    return {
-      message: { content: response.content[0].text, role: "assistant" },
-    };
+      const toolUseBlock = response.content.find(
+        (content): content is ToolUseBlock => content.type === "tool_use",
+      );
+
+      return {
+        raw: response,
+        message: {
+          content: response.content
+            .filter((content): content is TextBlock => content.type === "text")
+            .map((content) => ({
+              type: "text",
+              text: content.text,
+            })),
+          role: "assistant",
+          options: toolUseBlock
+            ? {
+                toolCall: {
+                  id: toolUseBlock.id,
+                  name: toolUseBlock.name,
+                  input: toolUseBlock.input,
+                },
+              }
+            : {},
+        },
+      };
+    } else {
+      const response = await anthropic.messages.create({
+        model: this.getModelName(this.model),
+        messages: this.formatMessages(messages),
+        max_tokens: this.maxTokens ?? 4096,
+        temperature: this.temperature,
+        top_p: this.topP,
+        ...(systemPrompt && { system: systemPrompt }),
+      });
+
+      return {
+        raw: response,
+        message: {
+          content: response.content[0].text,
+          role: "assistant",
+          options: {},
+        },
+      };
+    }
   }
 
   protected async *streamChat(
-    messages: ChatMessage[],
+    messages: ChatMessage<ToolCallLLMMessageOptions>[],
     systemPrompt?: string | null,
-  ): AsyncIterable<ChatResponseChunk> {
+  ): AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>> {
     const stream = await this.session.anthropic.messages.create({
       model: this.getModelName(this.model),
-      messages: this.formatMessages(messages),
+      messages: this.formatMessages<false>(messages),
       max_tokens: this.maxTokens ?? 4096,
       temperature: this.temperature,
       top_p: this.topP,
@@ -211,8 +342,27 @@ export class Anthropic extends BaseLLM {
       if (typeof content !== "string") continue;
 
       idx_counter++;
-      yield { delta: content };
+      yield {
+        raw: part,
+        delta: content,
+        options: {},
+      };
     }
     return;
   }
+
+  static toTool(tool: BaseTool): Tool {
+    if (tool.metadata.parameters?.type !== "object") {
+      throw new TypeError("Tool parameters must be an object");
+    }
+    return {
+      input_schema: {
+        type: "object",
+        properties: tool.metadata.parameters.properties,
+        required: tool.metadata.parameters.required,
+      },
+      name: tool.metadata.name,
+      description: tool.metadata.description,
+    };
+  }
 }

packages/core/src/llm/base.ts

@@ -8,14 +8,13 @@ import type {
   LLMCompletionParamsNonStreaming,
   LLMCompletionParamsStreaming,
   LLMMetadata,
+  ToolCallLLMMessageOptions,
 } from "./types.js";
 import { extractText, streamConverter } from "./utils.js";
 
 export abstract class BaseLLM<
-  AdditionalChatOptions extends Record<string, unknown> = Record<
-    string,
-    unknown
-  >,
+  AdditionalChatOptions extends object = object,
+  AdditionalMessageOptions extends object = object,
 > implements LLM<AdditionalChatOptions>
 {
   abstract metadata: LLMMetadata;
@@ -37,6 +36,7 @@ export abstract class BaseLLM<
       });
       return streamConverter(stream, (chunk) => {
         return {
+          raw: null,
           text: chunk.delta,
         };
       });
@@ -51,9 +51,21 @@ export abstract class BaseLLM<
   }
 
   abstract chat(
-    params: LLMChatParamsStreaming<AdditionalChatOptions>,
+    params: LLMChatParamsStreaming<
+      AdditionalChatOptions,
+      AdditionalMessageOptions
+    >,
   ): Promise<AsyncIterable<ChatResponseChunk>>;
   abstract chat(
-    params: LLMChatParamsNonStreaming<AdditionalChatOptions>,
-  ): Promise<ChatResponse>;
+    params: LLMChatParamsNonStreaming<
+      AdditionalChatOptions,
+      AdditionalMessageOptions
+    >,
+  ): Promise<ChatResponse<AdditionalMessageOptions>>;
+}
+
+export abstract class ToolCallLLM<
+  AdditionalChatOptions extends object = object,
+> extends BaseLLM<AdditionalChatOptions, ToolCallLLMMessageOptions> {
+  abstract supportToolCall: boolean;
 }

packages/core/src/llm/index.ts

@@ -1,4 +1,3 @@
-export * from "./LLM.js";
 export { Anthropic } from "./anthropic.js";
 export { FireworksLLM } from "./fireworks.js";
 export { Groq } from "./groq.js";
@@ -9,5 +8,12 @@ export {
 } from "./mistral.js";
 export { Ollama } from "./ollama.js";
 export * from "./open_ai.js";
+export { Portkey } from "./portkey.js";
+export * from "./replicate_ai.js";
+// Note: The type aliases for replicate are to simplify usage for Llama 2 (we're using replicate for Llama 2 support)
+export {
+  ReplicateChatStrategy as DeuceChatStrategy,
+  ReplicateLLM as LlamaDeuce,
+} from "./replicate_ai.js";
 export { TogetherLLM } from "./together.js";
 export * from "./types.js";

packages/core/src/llm/mistral.ts

@@ -106,6 +106,7 @@ export class MistralAI extends BaseLLM {
     const response = await client.chat(this.buildParams(messages));
     const message = response.choices[0].message;
     return {
+      raw: response,
       message,
     };
   }
@@ -136,6 +137,7 @@ export class MistralAI extends BaseLLM {
       idx_counter++;
 
       yield {
+        raw: part,
         delta: part.choices[0].delta.content ?? "",
       };
     }

packages/core/src/llm/ollama.ts

@@ -14,11 +14,13 @@ import type {
 
 const messageAccessor = (data: any): ChatResponseChunk => {
   return {
+    raw: data,
     delta: data.message.content,
   };
 };
+
 const completionAccessor = (data: any): CompletionResponse => {
-  return { text: data.response };
+  return { text: data.response, raw: null };
 };
 
 // https://github.com/jmorganca/ollama

packages/core/src/llm/open_ai.ts

@@ -7,7 +7,15 @@ import type {
 } from "openai";
 import { OpenAI as OrigOpenAI } from "openai";
 
-import type { ChatCompletionTool } from "openai/resources/chat/completions";
+import type {
+  ChatCompletionAssistantMessageParam,
+  ChatCompletionMessageToolCall,
+  ChatCompletionRole,
+  ChatCompletionSystemMessageParam,
+  ChatCompletionTool,
+  ChatCompletionToolMessageParam,
+  ChatCompletionUserMessageParam,
+} from "openai/resources/chat/completions";
 import type { ChatCompletionMessageParam } from "openai/resources/index.js";
 import { Tokenizers } from "../GlobalsHelper.js";
 import { wrapEventCaller } from "../internal/context/EventCaller.js";
@@ -20,7 +28,7 @@ import {
   getAzureModel,
   shouldUseAzure,
 } from "./azure.js";
-import { BaseLLM } from "./base.js";
+import { ToolCallLLM } from "./base.js";
 import type {
   ChatMessage,
   ChatResponse,
@@ -29,10 +37,11 @@ import type {
   LLMChatParamsNonStreaming,
   LLMChatParamsStreaming,
   LLMMetadata,
-  MessageToolCall,
   MessageType,
+  ToolCallLLMMessageOptions,
+  ToolCallOptions,
 } from "./types.js";
-import { wrapLLMEvent } from "./utils.js";
+import { extractText, wrapLLMEvent } from "./utils.js";
 
 export class AzureOpenAI extends OrigOpenAI {
   protected override authHeaders() {
@@ -135,6 +144,8 @@ export function isFunctionCallingModel(llm: LLM): llm is OpenAI {
   return isChatModel && !isOld;
 }
 
+export type OpenAIAdditionalMetadata = {};
+
 export type OpenAIAdditionalChatOptions = Omit<
   Partial<OpenAILLM.Chat.ChatCompletionCreateParams>,
   | "max_tokens"
@@ -147,11 +158,7 @@ export type OpenAIAdditionalChatOptions = Omit<
   | "toolChoice"
 >;
 
-export type OpenAIAdditionalMetadata = {
-  isFunctionCallingModel: boolean;
-};
-
-export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
+export class OpenAI extends ToolCallLLM<OpenAIAdditionalChatOptions> {
   // Per completion OpenAI params
   model: keyof typeof ALL_AVAILABLE_OPENAI_MODELS | string;
   temperature: number;
@@ -222,6 +229,10 @@ export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
     }
   }
 
+  get supportToolCall() {
+    return isFunctionCallingModel(this);
+  }
+
   get metadata(): LLMMetadata & OpenAIAdditionalMetadata {
     const contextWindow =
       ALL_AVAILABLE_OPENAI_MODELS[
@@ -234,13 +245,10 @@ export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
       maxTokens: this.maxTokens,
       contextWindow,
       tokenizer: Tokenizers.CL100K_BASE,
-      isFunctionCallingModel: isFunctionCallingModel(this),
     };
   }
 
-  mapMessageType(
-    messageType: MessageType,
-  ): "user" | "assistant" | "system" | "function" | "tool" {
+  static toOpenAIRole(messageType: MessageType): ChatCompletionRole {
     switch (messageType) {
       case "user":
         return "user";
@@ -248,52 +256,95 @@ export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
         return "assistant";
       case "system":
         return "system";
-      case "function":
-        return "function";
-      case "tool":
-        return "tool";
       default:
         return "user";
     }
   }
 
-  toOpenAIMessage(messages: ChatMessage[]) {
+  static toOpenAIMessage(
+    messages: ChatMessage<ToolCallLLMMessageOptions>[],
+  ): ChatCompletionMessageParam[] {
     return messages.map((message) => {
-      const additionalKwargs = message.additionalKwargs ?? {};
-
-      if (message.additionalKwargs?.toolCalls) {
-        additionalKwargs.tool_calls = message.additionalKwargs.toolCalls;
-        delete additionalKwargs.toolCalls;
+      const options = message.options ?? {};
+      if ("toolResult" in options) {
+        return {
+          tool_call_id: options.toolResult.id,
+          role: "tool",
+          content: extractText(message.content),
+        } satisfies ChatCompletionToolMessageParam;
+      } else if ("toolCall" in options) {
+        return {
+          role: "assistant",
+          content: extractText(message.content),
+          tool_calls: [
+            {
+              id: options.toolCall.id,
+              type: "function",
+              function: {
+                name: options.toolCall.name,
+                arguments:
+                  typeof options.toolCall.input === "string"
+                    ? options.toolCall.input
+                    : JSON.stringify(options.toolCall.input),
+              },
+            },
+          ],
+        } satisfies ChatCompletionAssistantMessageParam;
+      } else if (message.role === "user") {
+        return {
+          role: "user",
+          content: message.content,
+        } satisfies ChatCompletionUserMessageParam;
       }
 
-      return {
-        role: this.mapMessageType(message.role),
-        content: message.content,
-        ...additionalKwargs,
+      const response:
+        | ChatCompletionSystemMessageParam
+        | ChatCompletionUserMessageParam
+        | ChatCompletionMessageToolCall = {
+        // fixme(alex): type assertion
+        role: OpenAI.toOpenAIRole(message.role) as never,
+        // fixme: should not extract text, but assert content is string
+        content: extractText(message.content),
       };
+      return response;
     });
   }
 
   chat(
-    params: LLMChatParamsStreaming<OpenAIAdditionalChatOptions>,
-  ): Promise<AsyncIterable<ChatResponseChunk>>;
+    params: LLMChatParamsStreaming<
+      OpenAIAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >,
+  ): Promise<AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>>>;
   chat(
-    params: LLMChatParamsNonStreaming<OpenAIAdditionalChatOptions>,
-  ): Promise<ChatResponse>;
+    params: LLMChatParamsNonStreaming<
+      OpenAIAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >,
+  ): Promise<ChatResponse<ToolCallLLMMessageOptions>>;
   @wrapEventCaller
   @wrapLLMEvent
   async chat(
     params:
-      | LLMChatParamsNonStreaming<OpenAIAdditionalChatOptions>
-      | LLMChatParamsStreaming<OpenAIAdditionalChatOptions>,
-  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+      | LLMChatParamsNonStreaming<
+          OpenAIAdditionalChatOptions,
+          ToolCallLLMMessageOptions
+        >
+      | LLMChatParamsStreaming<
+          OpenAIAdditionalChatOptions,
+          ToolCallLLMMessageOptions
+        >,
+  ): Promise<
+    | ChatResponse<ToolCallLLMMessageOptions>
+    | AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>>
+  > {
     const { messages, stream, tools, additionalChatOptions } = params;
     const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
       model: this.model,
       temperature: this.temperature,
       max_tokens: this.maxTokens,
       tools: tools?.map(OpenAI.toTool),
-      messages: this.toOpenAIMessage(messages) as ChatCompletionMessageParam[],
+      messages: OpenAI.toOpenAIMessage(messages),
       top_p: this.topP,
       ...Object.assign({}, this.additionalChatOptions, additionalChatOptions),
     };
@@ -311,17 +362,21 @@ export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
 
     const content = response.choices[0].message?.content ?? "";
 
-    const kwargsOutput: Record<string, any> = {};
-
-    if (response.choices[0].message?.tool_calls) {
-      kwargsOutput.toolCalls = response.choices[0].message.tool_calls;
-    }
-
     return {
+      raw: response,
       message: {
         content,
         role: response.choices[0].message.role,
-        additionalKwargs: kwargsOutput,
+        options: response.choices[0].message?.tool_calls
+          ? {
+              toolCall: {
+                id: response.choices[0].message.tool_calls[0].id,
+                name: response.choices[0].message.tool_calls[0].function.name,
+                input:
+                  response.choices[0].message.tool_calls[0].function.arguments,
+              },
+            }
+          : {},
       },
     };
   }
@@ -329,7 +384,7 @@ export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
   @wrapEventCaller
   protected async *streamChat(
     baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams,
-  ): AsyncIterable<ChatResponseChunk> {
+  ): AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>> {
     const stream: AsyncIterable<OpenAILLM.Chat.ChatCompletionChunk> =
       await this.session.openai.chat.completions.create({
         ...baseRequestParams,
@@ -339,11 +394,26 @@ export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
     // TODO: add callback to streamConverter and use streamConverter here
     //Indices
     let idxCounter: number = 0;
-    const toolCalls: MessageToolCall[] = [];
+    let toolCallOptions: ToolCallOptions | null = null;
     for await (const part of stream) {
       if (!part.choices.length) continue;
       const choice = part.choices[0];
-      updateToolCalls(toolCalls, choice.delta.tool_calls);
+      // skip parts that don't have any content
+      if (!(choice.delta.content || choice.delta.tool_calls)) continue;
+      if (choice.delta.tool_calls?.[0].id) {
+        toolCallOptions = {
+          toolCall: {
+            name: choice.delta.tool_calls[0].function!.name!,
+            id: choice.delta.tool_calls[0].id,
+            input: choice.delta.tool_calls[0].function!.arguments!,
+          },
+        };
+      } else {
+        if (choice.delta.tool_calls?.[0].function?.arguments) {
+          toolCallOptions!.toolCall.input +=
+            choice.delta.tool_calls[0].function.arguments;
+        }
+      }
 
       const isDone: boolean = choice.finish_reason !== null;
 
@@ -354,9 +424,8 @@ export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
       });
 
       yield {
-        // add tool calls to final chunk
-        additionalKwargs:
-          toolCalls.length > 0 ? { toolCalls: toolCalls } : undefined,
+        raw: part,
+        options: toolCallOptions ? toolCallOptions : {},
         delta: choice.delta.content ?? "",
       };
     }
@@ -374,31 +443,3 @@ export class OpenAI extends BaseLLM<OpenAIAdditionalChatOptions> {
     };
   }
 }
-
-function updateToolCalls(
-  toolCalls: MessageToolCall[],
-  toolCallDeltas?: OpenAILLM.Chat.Completions.ChatCompletionChunk.Choice.Delta.ToolCall[],
-) {
-  function augmentToolCall(
-    toolCall?: MessageToolCall,
-    toolCallDelta?: OpenAILLM.Chat.Completions.ChatCompletionChunk.Choice.Delta.ToolCall,
-  ) {
-    toolCall =
-      toolCall ??
-      ({ function: { name: "", arguments: "" } } as MessageToolCall);
-    toolCall.id = toolCall.id ?? toolCallDelta?.id;
-    toolCall.type = toolCall.type ?? toolCallDelta?.type;
-    if (toolCallDelta?.function?.arguments) {
-      toolCall.function.arguments += toolCallDelta.function.arguments;
-    }
-    if (toolCallDelta?.function?.name) {
-      toolCall.function.name += toolCallDelta.function.name;
-    }
-    return toolCall;
-  }
-  if (toolCallDeltas) {
-    toolCallDeltas?.forEach((toolCall, i) => {
-      toolCalls[i] = augmentToolCall(toolCalls[i], toolCall);
-    });
-  }
-}

packages/core/src/llm/portkey.ts

@@ -1,7 +1,20 @@
 import { getEnv } from "@llamaindex/env";
 import _ from "lodash";
 import type { LLMOptions } from "portkey-ai";
-import { Portkey } from "portkey-ai";
+import { Portkey as OrigPortKey } from "portkey-ai";
+import { type StreamCallbackResponse } from "../callbacks/CallbackManager.js";
+import { getCallbackManager } from "../internal/settings/CallbackManager.js";
+import { BaseLLM } from "./base.js";
+import type {
+  ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
+  LLMChatParamsNonStreaming,
+  LLMChatParamsStreaming,
+  LLMMetadata,
+  MessageType,
+} from "./types.js";
+import { extractText, wrapLLMEvent } from "./utils.js";
 
 interface PortkeyOptions {
   apiKey?: string;
@@ -11,7 +24,7 @@ interface PortkeyOptions {
 }
 
 export class PortkeySession {
-  portkey: Portkey;
+  portkey: OrigPortKey;
 
   constructor(options: PortkeyOptions = {}) {
     if (!options.apiKey) {
@@ -22,13 +35,13 @@ export class PortkeySession {
       options.baseURL = getEnv("PORTKEY_BASE_URL") ?? "https://api.portkey.ai";
     }
 
-    this.portkey = new Portkey({});
+    this.portkey = new OrigPortKey({});
     this.portkey.llms = [{}];
     if (!options.apiKey) {
       throw new Error("Set Portkey ApiKey in PORTKEY_API_KEY env variable");
     }
 
-    this.portkey = new Portkey(options);
+    this.portkey = new OrigPortKey(options);
   }
 }
 
@@ -54,3 +67,92 @@ export function getPortkeySession(options: PortkeyOptions = {}) {
   }
   return session;
 }
+
+export class Portkey extends BaseLLM {
+  apiKey?: string = undefined;
+  baseURL?: string = undefined;
+  mode?: string = undefined;
+  llms?: [LLMOptions] | null = undefined;
+  session: PortkeySession;
+
+  constructor(init?: Partial<Portkey>) {
+    super();
+    this.apiKey = init?.apiKey;
+    this.baseURL = init?.baseURL;
+    this.mode = init?.mode;
+    this.llms = init?.llms;
+    this.session = getPortkeySession({
+      apiKey: this.apiKey,
+      baseURL: this.baseURL,
+      llms: this.llms,
+      mode: this.mode,
+    });
+  }
+
+  get metadata(): LLMMetadata {
+    throw new Error("metadata not implemented for Portkey");
+  }
+
+  chat(
+    params: LLMChatParamsStreaming,
+  ): Promise<AsyncIterable<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  @wrapLLMEvent
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    const { messages, stream, additionalChatOptions } = params;
+    if (stream) {
+      return this.streamChat(messages, additionalChatOptions);
+    } else {
+      const bodyParams = additionalChatOptions || {};
+      const response = await this.session.portkey.chatCompletions.create({
+        messages: messages.map((message) => ({
+          content: extractText(message.content),
+          role: message.role,
+        })),
+        ...bodyParams,
+      });
+
+      const content = response.choices[0].message?.content ?? "";
+      const role = response.choices[0].message?.role || "assistant";
+      return { raw: response, message: { content, role: role as MessageType } };
+    }
+  }
+
+  async *streamChat(
+    messages: ChatMessage[],
+    params?: Record<string, any>,
+  ): AsyncIterable<ChatResponseChunk> {
+    const chunkStream = await this.session.portkey.chatCompletions.create({
+      messages: messages.map((message) => ({
+        content: extractText(message.content),
+        role: message.role,
+      })),
+      ...params,
+      stream: true,
+    });
+
+    //Indices
+    let idx_counter: number = 0;
+    for await (const part of chunkStream) {
+      //Increment
+      part.choices[0].index = idx_counter;
+      const is_done: boolean =
+        part.choices[0].finish_reason === "stop" ? true : false;
+      //onLLMStream Callback
+
+      const stream_callback: StreamCallbackResponse = {
+        index: idx_counter,
+        isDone: is_done,
+        // token: part,
+      };
+      getCallbackManager().dispatchEvent("stream", stream_callback);
+
+      idx_counter++;
+
+      yield { raw: part, delta: part.choices[0].delta?.content ?? "" };
+    }
+    return;
+  }
+}

packages/core/src/llm/replicate_ai.ts

@@ -1,5 +1,15 @@
 import { getEnv } from "@llamaindex/env";
 import Replicate from "replicate";
+import { BaseLLM } from "./base.js";
+import type {
+  ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
+  LLMChatParamsNonStreaming,
+  LLMChatParamsStreaming,
+  MessageType,
+} from "./types.js";
+import { extractText, wrapLLMEvent } from "./utils.js";
 
 export class ReplicateSession {
   replicateKey: string | null = null;
@@ -20,12 +30,271 @@ export class ReplicateSession {
   }
 }
 
-let defaultReplicateSession: ReplicateSession | null = null;
+export const ALL_AVAILABLE_REPLICATE_MODELS = {
+  // TODO: add more models from replicate
+  "Llama-2-70b-chat-old": {
+    contextWindow: 4096,
+    replicateApi:
+      "replicate/llama70b-v2-chat:e951f18578850b652510200860fc4ea62b3b16fac280f83ff32282f87bbd2e48",
+    //^ Previous 70b model. This is also actually 4 bit, although not exllama.
+  },
+  "Llama-2-70b-chat-4bit": {
+    contextWindow: 4096,
+    replicateApi:
+      "meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3",
+    //^ Model is based off of exllama 4bit.
+  },
+  "Llama-2-13b-chat-old": {
+    contextWindow: 4096,
+    replicateApi:
+      "a16z-infra/llama13b-v2-chat:df7690f1994d94e96ad9d568eac121aecf50684a0b0963b25a41cc40061269e5",
+  },
+  //^ Last known good 13b non-quantized model. In future versions they add the SYS and INST tags themselves
+  "Llama-2-13b-chat-4bit": {
+    contextWindow: 4096,
+    replicateApi:
+      "meta/llama-2-13b-chat:f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
+  },
+  "Llama-2-7b-chat-old": {
+    contextWindow: 4096,
+    replicateApi:
+      "a16z-infra/llama7b-v2-chat:4f0a4744c7295c024a1de15e1a63c880d3da035fa1f49bfd344fe076074c8eea",
+    //^ Last (somewhat) known good 7b non-quantized model. In future versions they add the SYS and INST
+    // tags themselves
+    // https://github.com/replicate/cog-llama-template/commit/fa5ce83912cf82fc2b9c01a4e9dc9bff6f2ef137
+    // Problem is that they fix the max_new_tokens issue in the same commit. :-(
+  },
+  "Llama-2-7b-chat-4bit": {
+    contextWindow: 4096,
+    replicateApi:
+      "meta/llama-2-7b-chat:13c3cdee13ee059ab779f0291d29054dab00a47dad8261375654de5540165fb0",
+  },
+};
 
-export function getReplicateSession(replicateKey: string | null = null) {
-  if (!defaultReplicateSession) {
-    defaultReplicateSession = new ReplicateSession(replicateKey);
+export enum ReplicateChatStrategy {
+  A16Z = "a16z",
+  META = "meta",
+  METAWBOS = "metawbos",
+  //^ This is not exactly right because SentencePiece puts the BOS and EOS token IDs in after tokenization
+  // Unfortunately any string only API won't support these properly.
+  REPLICATE4BIT = "replicate4bit",
+  //^ To satisfy Replicate's 4 bit models' requirements where they also insert some INST tags
+  REPLICATE4BITWNEWLINES = "replicate4bitwnewlines",
+  //^ Replicate's documentation recommends using newlines: https://replicate.com/blog/how-to-prompt-llama
+}
+
+/**
+ * Replicate LLM implementation used
+ */
+export class ReplicateLLM extends BaseLLM {
+  model: keyof typeof ALL_AVAILABLE_REPLICATE_MODELS;
+  chatStrategy: ReplicateChatStrategy;
+  temperature: number;
+  topP: number;
+  maxTokens?: number;
+  replicateSession: ReplicateSession;
+
+  constructor(init?: Partial<ReplicateLLM>) {
+    super();
+    this.model = init?.model ?? "Llama-2-70b-chat-4bit";
+    this.chatStrategy =
+      init?.chatStrategy ??
+      (this.model.endsWith("4bit")
+        ? ReplicateChatStrategy.REPLICATE4BITWNEWLINES // With the newer Replicate models they do the system message themselves.
+        : ReplicateChatStrategy.METAWBOS); // With BOS and EOS seems to work best, although they all have problems past a certain point
+    this.temperature = init?.temperature ?? 0.1; // minimum temperature is 0.01 for Replicate endpoint
+    this.topP = init?.topP ?? 1;
+    this.maxTokens =
+      init?.maxTokens ??
+      ALL_AVAILABLE_REPLICATE_MODELS[this.model].contextWindow; // For Replicate, the default is 500 tokens which is too low.
+    this.replicateSession = init?.replicateSession ?? new ReplicateSession();
   }
 
-  return defaultReplicateSession;
+  get metadata() {
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      contextWindow: ALL_AVAILABLE_REPLICATE_MODELS[this.model].contextWindow,
+      tokenizer: undefined,
+    };
+  }
+
+  mapMessagesToPrompt(messages: ChatMessage[]) {
+    if (this.chatStrategy === ReplicateChatStrategy.A16Z) {
+      return this.mapMessagesToPromptA16Z(messages);
+    } else if (this.chatStrategy === ReplicateChatStrategy.META) {
+      return this.mapMessagesToPromptMeta(messages);
+    } else if (this.chatStrategy === ReplicateChatStrategy.METAWBOS) {
+      return this.mapMessagesToPromptMeta(messages, { withBos: true });
+    } else if (this.chatStrategy === ReplicateChatStrategy.REPLICATE4BIT) {
+      return this.mapMessagesToPromptMeta(messages, {
+        replicate4Bit: true,
+        withNewlines: true,
+      });
+    } else if (
+      this.chatStrategy === ReplicateChatStrategy.REPLICATE4BITWNEWLINES
+    ) {
+      return this.mapMessagesToPromptMeta(messages, {
+        replicate4Bit: true,
+        withNewlines: true,
+      });
+    } else {
+      return this.mapMessagesToPromptMeta(messages);
+    }
+  }
+
+  mapMessagesToPromptA16Z(messages: ChatMessage[]) {
+    return {
+      prompt:
+        messages.reduce((acc, message) => {
+          return (
+            (acc && `${acc}\n\n`) +
+            `${this.mapMessageTypeA16Z(message.role)}${message.content}`
+          );
+        }, "") + "\n\nAssistant:",
+      //^ Here we're differing from A16Z by omitting the space. Generally spaces at the end of prompts decrease performance due to tokenization
+      systemPrompt: undefined,
+    };
+  }
+
+  mapMessageTypeA16Z(messageType: MessageType): string {
+    switch (messageType) {
+      case "user":
+        return "User: ";
+      case "assistant":
+        return "Assistant: ";
+      case "system":
+        return "";
+      default:
+        throw new Error("Unsupported ReplicateLLM message type");
+    }
+  }
+
+  mapMessagesToPromptMeta(
+    messages: ChatMessage[],
+    opts?: {
+      withBos?: boolean;
+      replicate4Bit?: boolean;
+      withNewlines?: boolean;
+    },
+  ) {
+    const {
+      withBos = false,
+      replicate4Bit = false,
+      withNewlines = false,
+    } = opts ?? {};
+    const DEFAULT_SYSTEM_PROMPT = `You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.`;
+
+    const B_SYS = "<<SYS>>\n";
+    const E_SYS = "\n<</SYS>>\n\n";
+    const B_INST = "[INST]";
+    const E_INST = "[/INST]";
+    const BOS = "<s>";
+    const EOS = "</s>";
+
+    if (messages.length === 0) {
+      return { prompt: "", systemPrompt: undefined };
+    }
+
+    messages = [...messages]; // so we can use shift without mutating the original array
+
+    let systemPrompt = undefined;
+    if (messages[0].role === "system") {
+      const systemMessage = messages.shift()!;
+
+      if (replicate4Bit) {
+        systemPrompt = systemMessage.content;
+      } else {
+        const systemStr = `${B_SYS}${systemMessage.content}${E_SYS}`;
+
+        // TS Bug: https://github.com/microsoft/TypeScript/issues/9998
+        // @ts-ignore
+        if (messages[0].role !== "user") {
+          throw new Error(
+            "ReplicateLLM: if there is a system message, the second message must be a user message.",
+          );
+        }
+
+        const userContent = messages[0].content;
+
+        messages[0].content = `${systemStr}${userContent}`;
+      }
+    } else {
+      if (!replicate4Bit) {
+        messages[0].content = `${B_SYS}${DEFAULT_SYSTEM_PROMPT}${E_SYS}${messages[0].content}`;
+      }
+    }
+
+    return {
+      prompt: messages.reduce((acc, message, index) => {
+        const content = extractText(message.content);
+        if (index % 2 === 0) {
+          return (
+            `${acc}${withBos ? BOS : ""}${B_INST} ${content.trim()} ${E_INST}` +
+            (withNewlines ? "\n" : "")
+          );
+        } else {
+          return (
+            `${acc} ${content.trim()}` +
+            (withNewlines ? "\n" : " ") +
+            (withBos ? EOS : "")
+          ); // Yes, the EOS comes after the space. This is not a mistake.
+        }
+      }, ""),
+      systemPrompt,
+    };
+  }
+
+  chat(
+    params: LLMChatParamsStreaming,
+  ): Promise<AsyncIterable<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  @wrapLLMEvent
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    const { messages, stream } = params;
+    const api = ALL_AVAILABLE_REPLICATE_MODELS[this.model]
+      .replicateApi as `${string}/${string}:${string}`;
+
+    const { prompt, systemPrompt } = this.mapMessagesToPrompt(messages);
+
+    const replicateOptions: any = {
+      input: {
+        prompt,
+        system_prompt: systemPrompt,
+        temperature: this.temperature,
+        top_p: this.topP,
+      },
+    };
+
+    if (this.model.endsWith("4bit")) {
+      replicateOptions.input.max_new_tokens = this.maxTokens;
+    } else {
+      replicateOptions.input.max_length = this.maxTokens;
+    }
+
+    //TODO: Add streaming for this
+    if (stream) {
+      throw new Error("Streaming not supported for ReplicateLLM");
+    }
+
+    //Non-streaming
+    const response = await this.replicateSession.replicate.run(
+      api,
+      replicateOptions,
+    );
+    return {
+      raw: response,
+      message: {
+        content: (response as Array<string>).join("").trimStart(),
+        //^ We need to do this because Replicate returns a list of strings (for streaming functionality which is not exposed by the run function)
+        role: "assistant",
+      },
+    };
+  }
 }

packages/core/src/llm/types.ts

@@ -1,60 +1,68 @@
 import type { Tokenizers } from "../GlobalsHelper.js";
 import type { BaseTool, UUID } from "../types.js";
 
-type LLMBaseEvent<
-  Type extends string,
-  Payload extends Record<string, unknown>,
-> = CustomEvent<{
+type LLMBaseEvent<Payload extends Record<string, unknown>> = CustomEvent<{
   payload: Payload;
 }>;
 
-export type LLMStartEvent = LLMBaseEvent<
-  "llm-start",
-  {
-    id: UUID;
-    messages: ChatMessage[];
-  }
->;
-export type LLMEndEvent = LLMBaseEvent<
-  "llm-end",
-  {
-    id: UUID;
-    response: ChatResponse;
-  }
->;
+export type LLMStartEvent = LLMBaseEvent<{
+  id: UUID;
+  messages: ChatMessage[];
+}>;
+export type LLMToolCallEvent = LLMBaseEvent<{
+  // fixme: id is missing in the context
+  // id: UUID;
+  toolCall: Omit<ToolCallOptions["toolCall"], "id">;
+}>;
+export type LLMEndEvent = LLMBaseEvent<{
+  id: UUID;
+  response: ChatResponse;
+}>;
+export type LLMStreamEvent = LLMBaseEvent<{
+  id: UUID;
+  chunk: ChatResponseChunk;
+}>;
 
 /**
  * @internal
  */
 export interface LLMChat<
-  ExtraParams extends Record<string, unknown> = Record<string, unknown>,
+  AdditionalChatOptions extends object = object,
+  AdditionalMessageOptions extends object = object,
 > {
   chat(
     params:
-      | LLMChatParamsStreaming<ExtraParams>
-      | LLMChatParamsNonStreaming<ExtraParams>,
-  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>>;
+      | LLMChatParamsStreaming<AdditionalChatOptions>
+      | LLMChatParamsNonStreaming<AdditionalChatOptions>,
+  ): Promise<
+    | ChatResponse<AdditionalMessageOptions>
+    | AsyncIterable<ChatResponseChunk<AdditionalMessageOptions>>
+  >;
 }
 
 /**
  * Unified language model interface
  */
 export interface LLM<
-  AdditionalChatOptions extends Record<string, unknown> = Record<
-    string,
-    unknown
-  >,
+  AdditionalChatOptions extends object = object,
+  AdditionalMessageOptions extends object = object,
 > extends LLMChat<AdditionalChatOptions> {
   metadata: LLMMetadata;
   /**
    * Get a chat response from the LLM
    */
   chat(
-    params: LLMChatParamsStreaming<AdditionalChatOptions>,
+    params: LLMChatParamsStreaming<
+      AdditionalChatOptions,
+      AdditionalMessageOptions
+    >,
   ): Promise<AsyncIterable<ChatResponseChunk>>;
   chat(
-    params: LLMChatParamsNonStreaming<AdditionalChatOptions>,
-  ): Promise<ChatResponse>;
+    params: LLMChatParamsNonStreaming<
+      AdditionalChatOptions,
+      AdditionalMessageOptions
+    >,
+  ): Promise<ChatResponse<AdditionalMessageOptions>>;
 
   /**
    * Get a prompt completion from the LLM
@@ -67,35 +75,42 @@ export interface LLM<
   ): Promise<CompletionResponse>;
 }
 
-export type MessageType =
-  | "user"
-  | "assistant"
-  | "system"
-  | "generic"
-  | "function"
-  | "memory"
-  | "tool";
+export type MessageType = "user" | "assistant" | "system" | "memory";
 
-export interface ChatMessage {
+export type ChatMessage<AdditionalMessageOptions extends object = object> = {
   content: MessageContent;
   role: MessageType;
-  additionalKwargs?: Record<string, any>;
+  options?: undefined | AdditionalMessageOptions;
+};
+
+export interface ChatResponse<
+  AdditionalMessageOptions extends object = object,
+> {
+  message: ChatMessage<AdditionalMessageOptions>;
+  /**
+   * Raw response from the LLM
+   *
+   * If LLM response an iterable of chunks, this will be an array of those chunks
+   */
+  raw: object | null;
 }
 
-export interface ChatResponse {
-  message: ChatMessage;
-  raw?: Record<string, any>;
-  additionalKwargs?: Record<string, any>;
-}
-
-export interface ChatResponseChunk {
+export type ChatResponseChunk<
+  AdditionalMessageOptions extends object = object,
+> = {
+  raw: object | null;
   delta: string;
-  additionalKwargs?: Record<string, any>;
-}
+  options?: undefined | AdditionalMessageOptions;
+};
 
 export interface CompletionResponse {
   text: string;
-  raw?: Record<string, any>;
+  /**
+   * Raw response from the LLM
+   *
+   * It's possible that this is `null` if the LLM response an iterable of chunks
+   */
+  raw: object | null;
 }
 
 export type LLMMetadata = {
@@ -108,32 +123,25 @@ export type LLMMetadata = {
 };
 
 export interface LLMChatParamsBase<
-  AdditionalChatOptions extends Record<string, unknown> = Record<
-    string,
-    unknown
-  >,
+  AdditionalChatOptions extends object = object,
+  AdditionalMessageOptions extends object = object,
 > {
-  messages: ChatMessage[];
+  messages: ChatMessage<AdditionalMessageOptions>[];
   additionalChatOptions?: AdditionalChatOptions;
   tools?: BaseTool[];
-  additionalKwargs?: Record<string, unknown>;
 }
 
 export interface LLMChatParamsStreaming<
-  AdditionalChatOptions extends Record<string, unknown> = Record<
-    string,
-    unknown
-  >,
-> extends LLMChatParamsBase<AdditionalChatOptions> {
+  AdditionalChatOptions extends object = object,
+  AdditionalMessageOptions extends object = object,
+> extends LLMChatParamsBase<AdditionalChatOptions, AdditionalMessageOptions> {
   stream: true;
 }
 
 export interface LLMChatParamsNonStreaming<
-  AdditionalChatOptions extends Record<string, unknown> = Record<
-    string,
-    unknown
-  >,
-> extends LLMChatParamsBase<AdditionalChatOptions> {
+  AdditionalChatOptions extends object = object,
+  AdditionalMessageOptions extends object = object,
+> extends LLMChatParamsBase<AdditionalChatOptions, AdditionalMessageOptions> {
   stream?: false;
 }
 
@@ -169,13 +177,29 @@ export type MessageContentDetail =
  */
 export type MessageContent = string | MessageContentDetail[];
 
-interface Function {
-  arguments: string;
+export type ToolCall = {
   name: string;
-}
-
-export interface MessageToolCall {
+  // for now, claude-3-opus will give object, gpt-3/4 will give string
+  // todo: unify this to always be an object
+  input: unknown;
   id: string;
-  function: Function;
-  type: "function";
-}
+};
+
+export type ToolResult = {
+  id: string;
+  result: string;
+  isError: boolean;
+};
+
+export type ToolCallOptions = {
+  toolCall: ToolCall;
+};
+
+export type ToolResultOptions = {
+  toolResult: ToolResult;
+};
+
+export type ToolCallLLMMessageOptions =
+  | ToolResultOptions
+  | ToolCallOptions
+  | {};

packages/core/src/llm/utils.ts

@@ -2,6 +2,7 @@ import { AsyncLocalStorage, randomUUID } from "@llamaindex/env";
 import { getCallbackManager } from "../internal/settings/CallbackManager.js";
 import type {
   ChatResponse,
+  ChatResponseChunk,
   LLM,
   LLMChat,
   MessageContent,
@@ -60,14 +61,24 @@ export function extractText(message: MessageContent): string {
 /**
  * @internal
  */
-export function wrapLLMEvent(
-  originalMethod: LLMChat["chat"],
+export function wrapLLMEvent<
+  AdditionalChatOptions extends object = object,
+  AdditionalMessageOptions extends object = object,
+>(
+  originalMethod: LLMChat<
+    AdditionalChatOptions,
+    AdditionalMessageOptions
+  >["chat"],
   _context: ClassMethodDecoratorContext,
 ) {
   return async function withLLMEvent(
-    this: LLM,
-    ...params: Parameters<LLMChat["chat"]>
-  ): ReturnType<LLMChat["chat"]> {
+    this: LLM<AdditionalChatOptions, AdditionalMessageOptions>,
+    ...params: Parameters<
+      LLMChat<AdditionalChatOptions, AdditionalMessageOptions>["chat"]
+    >
+  ): ReturnType<
+    LLMChat<AdditionalChatOptions, AdditionalMessageOptions>["chat"]
+  > {
     const id = randomUUID();
     getCallbackManager().dispatchEvent("llm-start", {
       payload: {
@@ -83,12 +94,14 @@ export function wrapLLMEvent(
         [Symbol.asyncIterator]: response[Symbol.asyncIterator].bind(response),
       };
       response[Symbol.asyncIterator] = async function* () {
-        const finalResponse: ChatResponse = {
+        const finalResponse = {
+          raw: [] as ChatResponseChunk[],
           message: {
             content: "",
             role: "assistant",
+            options: {},
           },
-        };
+        } satisfies ChatResponse;
         let firstOne = false;
         for await (const chunk of originalAsyncIterator) {
           if (!firstOne) {
@@ -97,6 +110,19 @@ export function wrapLLMEvent(
           } else {
             finalResponse.message.content += chunk.delta;
           }
+          if (chunk.options) {
+            finalResponse.message.options = {
+              ...finalResponse.message.options,
+              ...chunk.options,
+            };
+          }
+          getCallbackManager().dispatchEvent("llm-stream", {
+            payload: {
+              id,
+              chunk,
+            },
+          });
+          finalResponse.raw.push(chunk);
           yield chunk;
         }
         snapshot(() => {

packages/core/src/memory/ChatMemoryBuffer.ts

@@ -1,3 +1,4 @@
+import type { ChatHistory } from "../ChatHistory.js";
 import type { ChatMessage, LLM } from "../llm/index.js";
 import { SimpleChatStore } from "../storage/chatStore/SimpleChatStore.js";
 import type { BaseChatStore } from "../storage/chatStore/types.js";
@@ -6,28 +7,28 @@ import type { BaseMemory } from "./types.js";
 const DEFAULT_TOKEN_LIMIT_RATIO = 0.75;
 const DEFAULT_TOKEN_LIMIT = 3000;
 
-type ChatMemoryBufferParams = {
-  tokenLimit?: number;
-  chatStore?: BaseChatStore;
-  chatStoreKey?: string;
-  chatHistory?: ChatMessage[];
-  llm?: LLM;
-};
+type ChatMemoryBufferParams<AdditionalMessageOptions extends object = object> =
+  {
+    tokenLimit?: number;
+    chatStore?: BaseChatStore<AdditionalMessageOptions>;
+    chatStoreKey?: string;
+    chatHistory?: ChatHistory<AdditionalMessageOptions>;
+    llm?: LLM<object, AdditionalMessageOptions>;
+  };
 
-/**
- * Chat memory buffer.
- */
-export class ChatMemoryBuffer implements BaseMemory {
+export class ChatMemoryBuffer<AdditionalMessageOptions extends object = object>
+  implements BaseMemory<AdditionalMessageOptions>
+{
   tokenLimit: number;
 
-  chatStore: BaseChatStore;
+  chatStore: BaseChatStore<AdditionalMessageOptions>;
   chatStoreKey: string;
 
-  /**
-   * Initialize.
-   */
-  constructor(init?: Partial<ChatMemoryBufferParams>) {
-    this.chatStore = init?.chatStore ?? new SimpleChatStore();
+  constructor(
+    init?: Partial<ChatMemoryBufferParams<AdditionalMessageOptions>>,
+  ) {
+    this.chatStore =
+      init?.chatStore ?? new SimpleChatStore<AdditionalMessageOptions>();
     this.chatStoreKey = init?.chatStoreKey ?? "chat_history";
     if (init?.llm) {
       const contextWindow = init.llm.metadata.contextWindow;
@@ -39,15 +40,11 @@ export class ChatMemoryBuffer implements BaseMemory {
     }
 
     if (init?.chatHistory) {
-      this.chatStore.setMessages(this.chatStoreKey, init.chatHistory);
+      this.chatStore.setMessages(this.chatStoreKey, init.chatHistory.messages);
     }
   }
 
-  /**
-    Get chat history.
-    @param initialTokenCount: number of tokens to start with
-  */
-  get(initialTokenCount: number = 0): ChatMessage[] {
+  get(initialTokenCount: number = 0) {
     const chatHistory = this.getAll();
 
     if (initialTokenCount > this.tokenLimit) {
@@ -79,42 +76,22 @@ export class ChatMemoryBuffer implements BaseMemory {
     return chatHistory.slice(-messageCount);
   }
 
-  /**
-   * Get all chat history.
-   * @returns {ChatMessage[]} chat history
-   */
-  getAll(): ChatMessage[] {
+  getAll() {
     return this.chatStore.getMessages(this.chatStoreKey);
   }
 
-  /**
-   * Put chat history.
-   * @param message
-   */
-  put(message: ChatMessage): void {
+  put(message: ChatMessage<AdditionalMessageOptions>) {
     this.chatStore.addMessage(this.chatStoreKey, message);
   }
 
-  /**
-   * Set chat history.
-   * @param messages
-   */
-  set(messages: ChatMessage[]): void {
+  set(messages: ChatMessage<AdditionalMessageOptions>[]) {
     this.chatStore.setMessages(this.chatStoreKey, messages);
   }
 
-  /**
-   * Reset chat history.
-   */
-  reset(): void {
+  reset() {
     this.chatStore.deleteMessages(this.chatStoreKey);
   }
 
-  /**
-   * Get token count for message count.
-   * @param messageCount
-   * @returns {number} token count
-   */
   private _tokenCountForMessageCount(messageCount: number): number {
     if (messageCount <= 0) {
       return 0;

packages/core/src/memory/types.ts

@@ -1,24 +1,10 @@
 import type { ChatMessage } from "../llm/index.js";
 
-export interface BaseMemory {
-  /*
-   Get chat history.
-  */
-  get(...args: any): ChatMessage[];
-  /*
-    Get all chat history.
-  */
-  getAll(): ChatMessage[];
-  /*
-    Put chat history.
-  */
-  put(message: ChatMessage): void;
-  /*
-    Set chat history.
-  */
-  set(messages: ChatMessage[]): void;
-  /*
-    Reset chat history.
-  */
+export interface BaseMemory<AdditionalMessageOptions extends object = object> {
+  tokenLimit: number;
+  get(...args: unknown[]): ChatMessage<AdditionalMessageOptions>[];
+  getAll(): ChatMessage<AdditionalMessageOptions>[];
+  put(message: ChatMessage<AdditionalMessageOptions>): void;
+  set(messages: ChatMessage<AdditionalMessageOptions>[]): void;
   reset(): void;
 }

packages/core/src/objects/base.ts

@@ -50,7 +50,7 @@ export abstract class BaseObjectNodeMapping {
 
 type QueryType = string;
 
-export class ObjectRetriever {
+export class ObjectRetriever<T = unknown> {
   _retriever: BaseRetriever;
   _objectNodeMapping: BaseObjectNodeMapping;
 
@@ -68,7 +68,7 @@ export class ObjectRetriever {
   }
 
   // Translating the retrieve method
-  async retrieve(strOrQueryBundle: QueryType): Promise<any> {
+  async retrieve(strOrQueryBundle: QueryType): Promise<T[]> {
     const nodes = await this.retriever.retrieve({ query: strOrQueryBundle });
     const objs = nodes.map((n) => this._objectNodeMapping.fromNode(n.node));
     return objs;
@@ -180,7 +180,7 @@ export class ObjectIndex {
     return this._objectNodeMapping.objNodeMapping();
   }
 
-  async asRetriever(kwargs: any): Promise<ObjectRetriever> {
+  async asRetriever(kwargs: any): Promise<ObjectRetriever<any>> {
     return new ObjectRetriever(
       this._index.asRetriever(kwargs),
       this._objectNodeMapping,

packages/core/src/storage/chatStore/SimpleChatStore.ts

@@ -2,47 +2,35 @@ import type { ChatMessage } from "../../llm/index.js";
 import type { BaseChatStore } from "./types.js";
 
 /**
- * Simple chat store.
+ * fixme: User could carry object references in the messages.
+ *  This could lead to memory leaks if the messages are not properly cleaned up.
  */
-export class SimpleChatStore implements BaseChatStore {
-  store: { [key: string]: ChatMessage[] } = {};
+export class SimpleChatStore<
+  AdditionalMessageOptions extends object = Record<string, unknown>,
+> implements BaseChatStore<AdditionalMessageOptions>
+{
+  store: { [key: string]: ChatMessage<AdditionalMessageOptions>[] } = {};
 
-  /**
-   * Set messages.
-   * @param key: key
-   * @param messages: messages
-   * @returns: void
-   */
-  public setMessages(key: string, messages: ChatMessage[]): void {
+  public setMessages(
+    key: string,
+    messages: ChatMessage<AdditionalMessageOptions>[],
+  ) {
     this.store[key] = messages;
   }
 
-  /**
-   * Get messages.
-   * @param key: key
-   * @returns: messages
-   */
-  public getMessages(key: string): ChatMessage[] {
+  public getMessages(key: string): ChatMessage<AdditionalMessageOptions>[] {
     return this.store[key] || [];
   }
 
-  /**
-   * Add message.
-   * @param key: key
-   * @param message: message
-   * @returns: void
-   */
-  public addMessage(key: string, message: ChatMessage): void {
+  public addMessage(
+    key: string,
+    message: ChatMessage<AdditionalMessageOptions>,
+  ) {
     this.store[key] = this.store[key] || [];
     this.store[key].push(message);
   }
 
-  /**
-   * Delete messages.
-   * @param key: key
-   * @returns: messages
-   */
-  public deleteMessages(key: string): ChatMessage[] | null {
+  public deleteMessages(key: string) {
     if (!(key in this.store)) {
       return null;
     }
@@ -51,13 +39,7 @@ export class SimpleChatStore implements BaseChatStore {
     return messages;
   }
 
-  /**
-   * Delete message.
-   * @param key: key
-   * @param idx: idx
-   * @returns: message
-   */
-  public deleteMessage(key: string, idx: number): ChatMessage | null {
+  public deleteMessage(key: string, idx: number) {
     if (!(key in this.store)) {
       return null;
     }
@@ -67,12 +49,7 @@ export class SimpleChatStore implements BaseChatStore {
     return this.store[key].splice(idx, 1)[0];
   }
 
-  /**
-   * Delete last message.
-   * @param key: key
-   * @returns: message
-   */
-  public deleteLastMessage(key: string): ChatMessage | null {
+  public deleteLastMessage(key: string) {
     if (!(key in this.store)) {
       return null;
     }
@@ -82,10 +59,6 @@ export class SimpleChatStore implements BaseChatStore {
     return lastMessage || null;
   }
 
-  /**
-   * Get keys.
-   * @returns: keys
-   */
   public getKeys(): string[] {
     return Object.keys(this.store);
   }

packages/core/src/storage/chatStore/types.ts

@@ -1,11 +1,19 @@
 import type { ChatMessage } from "../../llm/index.js";
 
-export interface BaseChatStore {
-  setMessages(key: string, messages: ChatMessage[]): void;
-  getMessages(key: string): ChatMessage[];
-  addMessage(key: string, message: ChatMessage): void;
-  deleteMessages(key: string): ChatMessage[] | null;
-  deleteMessage(key: string, idx: number): ChatMessage | null;
-  deleteLastMessage(key: string): ChatMessage | null;
+export interface BaseChatStore<
+  AdditionalMessageOptions extends object = object,
+> {
+  setMessages(
+    key: string,
+    messages: ChatMessage<AdditionalMessageOptions>[],
+  ): void;
+  getMessages(key: string): ChatMessage<AdditionalMessageOptions>[];
+  addMessage(key: string, message: ChatMessage<AdditionalMessageOptions>): void;
+  deleteMessages(key: string): ChatMessage<AdditionalMessageOptions>[] | null;
+  deleteMessage(
+    key: string,
+    idx: number,
+  ): ChatMessage<AdditionalMessageOptions> | null;
+  deleteLastMessage(key: string): ChatMessage<AdditionalMessageOptions> | null;
   getKeys(): string[];
 }

packages/core/src/synthesizers/MultiModalResponseSynthesizer.ts

@@ -91,6 +91,6 @@ export class MultiModalResponseSynthesizer
       prompt,
     });
 
-    return new Response(response.text, nodes);
+    return new Response(response.text, nodesWithScore);
   }
 }

packages/core/src/synthesizers/ResponseSynthesizer.ts

@@ -69,19 +69,21 @@ export class ResponseSynthesizer
     const textChunks: string[] = nodesWithScore.map(({ node }) =>
       node.getContent(this.metadataMode),
     );
-    const nodes = nodesWithScore.map(({ node }) => node);
     if (stream) {
       const response = await this.responseBuilder.getResponse({
         query,
         textChunks,
         stream,
       });
-      return streamConverter(response, (chunk) => new Response(chunk, nodes));
+      return streamConverter(
+        response,
+        (chunk) => new Response(chunk, nodesWithScore),
+      );
     }
     const response = await this.responseBuilder.getResponse({
       query,
       textChunks,
     });
-    return new Response(response, nodes);
+    return new Response(response, nodesWithScore);
   }
 }

packages/core/src/tools/QueryEngineTool.ts

@@ -1,18 +1,11 @@
+import type { JSONSchemaType } from "ajv";
 import type { BaseQueryEngine, BaseTool, ToolMetadata } from "../types.js";
 
-export type QueryEngineToolParams = {
-  queryEngine: BaseQueryEngine;
-  metadata: ToolMetadata;
-};
-
-type QueryEngineCallParams = {
-  query: string;
-};
-
 const DEFAULT_NAME = "query_engine_tool";
 const DEFAULT_DESCRIPTION =
   "Useful for running a natural language query against a knowledge base and get back a natural language response.";
-const DEFAULT_PARAMETERS = {
+
+const DEFAULT_PARAMETERS: JSONSchemaType<QueryEngineParam> = {
   type: "object",
   properties: {
     query: {
@@ -23,9 +16,18 @@ const DEFAULT_PARAMETERS = {
   required: ["query"],
 };
 
-export class QueryEngineTool implements BaseTool {
+export type QueryEngineToolParams = {
+  queryEngine: BaseQueryEngine;
+  metadata: ToolMetadata<JSONSchemaType<QueryEngineParam>>;
+};
+
+export type QueryEngineParam = {
+  query: string;
+};
+
+export class QueryEngineTool implements BaseTool<QueryEngineParam> {
   private queryEngine: BaseQueryEngine;
-  metadata: ToolMetadata;
+  metadata: ToolMetadata<JSONSchemaType<QueryEngineParam>>;
 
   constructor({ queryEngine, metadata }: QueryEngineToolParams) {
     this.queryEngine = queryEngine;
@@ -36,18 +38,8 @@ export class QueryEngineTool implements BaseTool {
     };
   }
 
-  async call(...args: QueryEngineCallParams[]): Promise<any> {
-    let queryStr: string;
-
-    if (args && args.length > 0) {
-      queryStr = String(args[0].query);
-    } else {
-      throw new Error(
-        "Cannot call query engine without specifying `input` parameter.",
-      );
-    }
-
-    const response = await this.queryEngine.query({ query: queryStr });
+  async call({ query }: QueryEngineParam) {
+    const response = await this.queryEngine.query({ query });
 
     return response.response;
   }

packages/core/src/tools/ToolFactory.ts

@@ -1,33 +0,0 @@
-import type { BaseTool } from "../types.js";
-import { WikipediaTool } from "./WikipediaTool.js";
-
-enum Tools {
-  Wikipedia = "wikipedia.WikipediaToolSpec",
-}
-
-type ToolConfig = { [key in Tools]: Record<string, any> };
-
-export class ToolFactory {
-  private static async createTool(
-    key: Tools,
-    options: Record<string, any>,
-  ): Promise<BaseTool> {
-    if (key === Tools.Wikipedia) {
-      const tool = new WikipediaTool();
-      return tool;
-    }
-
-    throw new Error(
-      `Sorry! Tool ${key} is not supported yet. Options: ${options}`,
-    );
-  }
-
-  public static async createTools(config: ToolConfig): Promise<BaseTool[]> {
-    const tools: BaseTool[] = [];
-    for (const [key, value] of Object.entries(config as ToolConfig)) {
-      const tool = await ToolFactory.createTool(key as Tools, value);
-      tools.push(tool);
-    }
-    return tools;
-  }
-}

packages/core/src/tools/ToolsFactory.ts

@@ -0,0 +1,43 @@
+import { WikipediaTool } from "./WikipediaTool.js";
+
+export namespace ToolsFactory {
+  type ToolsMap = {
+    [Tools.Wikipedia]: typeof WikipediaTool;
+  };
+
+  export enum Tools {
+    Wikipedia = "wikipedia.WikipediaToolSpec",
+  }
+
+  export async function createTool<Tool extends Tools>(
+    key: Tool,
+    ...params: ConstructorParameters<ToolsMap[Tool]>
+  ): Promise<InstanceType<ToolsMap[Tool]>> {
+    if (key === Tools.Wikipedia) {
+      return new WikipediaTool(...params) as InstanceType<ToolsMap[Tool]>;
+    }
+
+    throw new Error(
+      `Sorry! Tool ${key} is not supported yet. Options: ${params}`,
+    );
+  }
+
+  export async function createTools<const Tool extends Tools>(record: {
+    [key in Tool]: ConstructorParameters<ToolsMap[Tool]>[1] extends any // backward compatibility for `create-llama` script // if parameters are an array, use them as is
+      ? ConstructorParameters<ToolsMap[Tool]>[0]
+      : ConstructorParameters<ToolsMap[Tool]>;
+  }): Promise<InstanceType<ToolsMap[Tool]>[]> {
+    const tools: InstanceType<ToolsMap[Tool]>[] = [];
+    for (const key in record) {
+      const params = record[key];
+      tools.push(
+        await createTool(
+          key,
+          // @ts-expect-error
+          Array.isArray(params) ? params : [params],
+        ),
+      );
+    }
+    return tools;
+  }
+}

packages/core/src/tools/WikipediaTool.ts

@@ -1,16 +1,17 @@
+import type { JSONSchemaType } from "ajv";
 import { default as wiki } from "wikipedia";
 import type { BaseTool, ToolMetadata } from "../types.js";
 
-export type WikipediaToolParams = {
-  metadata?: ToolMetadata;
-};
-
-type WikipediaCallParams = {
+type WikipediaParameter = {
   query: string;
   lang?: string;
 };
 
-const DEFAULT_META_DATA: ToolMetadata = {
+export type WikipediaToolParams = {
+  metadata?: ToolMetadata<JSONSchemaType<WikipediaParameter>>;
+};
+
+const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<WikipediaParameter>> = {
   name: "wikipedia_tool",
   description: "A tool that uses a query engine to search Wikipedia.",
   parameters: {
@@ -20,14 +21,19 @@ const DEFAULT_META_DATA: ToolMetadata = {
         type: "string",
         description: "The query to search for",
       },
+      lang: {
+        type: "string",
+        description: "The language to search in",
+        nullable: true,
+      },
     },
     required: ["query"],
   },
 };
 
-export class WikipediaTool implements BaseTool {
+export class WikipediaTool implements BaseTool<WikipediaParameter> {
   private readonly DEFAULT_LANG = "en";
-  metadata: ToolMetadata;
+  metadata: ToolMetadata<JSONSchemaType<WikipediaParameter>>;
 
   constructor(params?: WikipediaToolParams) {
     this.metadata = params?.metadata || DEFAULT_META_DATA;
@@ -46,7 +52,7 @@ export class WikipediaTool implements BaseTool {
   async call({
     query,
     lang = this.DEFAULT_LANG,
-  }: WikipediaCallParams): Promise<string> {
+  }: WikipediaParameter): Promise<string> {
     const searchResult = await wiki.default.search(query);
     if (searchResult.results.length === 0) return "No search results.";
     return await this.loadData(searchResult.results[0].title, lang);

packages/core/src/tools/functionTool.ts

@@ -1,32 +1,30 @@
+import type { JSONSchemaType } from "ajv";
 import type { BaseTool, ToolMetadata } from "../types.js";
 
-type Metadata = {
-  name: string;
-  description: string;
-  parameters: ToolMetadata["parameters"];
-};
+export class FunctionTool<T, R extends string | Promise<string>>
+  implements BaseTool<T>
+{
+  constructor(
+    private readonly _fn: (input: T) => R,
+    private readonly _metadata: ToolMetadata<JSONSchemaType<T>>,
+  ) {}
 
-export class FunctionTool<T = any> implements BaseTool {
-  private _fn: (...args: any[]) => any;
-  private _metadata: ToolMetadata;
-
-  constructor(fn: (...args: any[]) => any, metadata: Metadata) {
-    this._fn = fn;
-    this._metadata = metadata as ToolMetadata;
+  static from<T>(
+    fn: (input: T) => string | Promise<string>,
+    schema: ToolMetadata<JSONSchemaType<T>>,
+  ): FunctionTool<T, string | Promise<string>>;
+  static from<T, R extends string | Promise<string>>(
+    fn: (input: T) => R,
+    schema: ToolMetadata<JSONSchemaType<T>>,
+  ): FunctionTool<T, R> {
+    return new FunctionTool(fn, schema);
   }
 
-  static fromDefaults<T = any>(
-    fn: (...args: any[]) => any,
-    metadata?: Metadata,
-  ): FunctionTool<T> {
-    return new FunctionTool(fn, metadata!);
+  get metadata(): BaseTool<T>["metadata"] {
+    return this._metadata as BaseTool<T>["metadata"];
   }
 
-  get metadata(): ToolMetadata {
-    return this._metadata;
-  }
-
-  async call(...args: any[]): Promise<any> {
-    return this._fn(...args);
+  call(input: T) {
+    return this._fn(input);
   }
 }

packages/core/src/tools/index.ts

@@ -1,5 +1,4 @@
 export * from "./QueryEngineTool.js";
-export * from "./ToolFactory.js";
 export * from "./WikipediaTool.js";
 export * from "./functionTool.js";
 export * from "./types.js";

packages/core/src/tools/utils.ts

@@ -1,33 +1,33 @@
+import { getCallbackManager } from "../internal/settings/CallbackManager.js";
 import type { BaseTool } from "../types.js";
 import { ToolOutput } from "./types.js";
 
-/**
- * Call tool with error handling.
- * @param tool: tool
- * @param inputDict: input dict
- * @param errorMessage: error message
- * @param raiseError: raise error
- * @returns: tool output
- */
 export async function callToolWithErrorHandling(
   tool: BaseTool,
-  inputDict: { [key: string]: any },
-  errorMessage: string | null = null,
-  raiseError: boolean = false,
+  input: unknown,
 ): Promise<ToolOutput> {
-  try {
-    const value = await tool.call?.(inputDict);
-    return new ToolOutput(value, tool.metadata.name, inputDict, value);
-  } catch (e) {
-    if (raiseError) {
-      throw e;
-    }
-    errorMessage = errorMessage || `Error: ${e}`;
+  if (!tool.call) {
     return new ToolOutput(
-      errorMessage,
+      "Error: Tool does not have a call function.",
       tool.metadata.name,
-      { kwargs: inputDict },
-      e,
+      input,
+      null,
     );
   }
+  try {
+    getCallbackManager().dispatchEvent("llm-tool-call", {
+      payload: {
+        toolCall: {
+          name: tool.metadata.name,
+          input,
+        },
+      },
+    });
+    const value = await tool.call(
+      typeof input === "string" ? JSON.parse(input) : input,
+    );
+    return new ToolOutput(value, tool.metadata.name, input, value);
+  } catch (e) {
+    return new ToolOutput(`Error: ${e}`, tool.metadata.name, input, e);
+  }
 }

packages/core/src/types.ts

@@ -1,6 +1,7 @@
 /**
  * Top level types to avoid circular dependencies
  */
+import { type JSONSchemaType } from "ajv";
 import type { Response } from "./Response.js";
 
 /**
@@ -30,14 +31,46 @@ export interface BaseQueryEngine {
   query(params: QueryEngineParamsNonStreaming): Promise<Response>;
 }
 
+type Known =
+  | { [key: string]: Known }
+  | [Known, ...Known[]]
+  | Known[]
+  | number
+  | string
+  | boolean
+  | null;
+
+export type ToolMetadata<
+  Parameters extends Record<string, unknown> = Record<string, unknown>,
+> = {
+  description: string;
+  name: string;
+  /**
+   * OpenAI uses JSON Schema to describe the parameters that a tool can take.
+   * @link https://json-schema.org/understanding-json-schema
+   */
+  parameters?: Parameters;
+};
+
 /**
  * Simple Tool interface. Likely to change.
  */
-export interface BaseTool {
-  call?: (...args: any[]) => any;
-  metadata: ToolMetadata;
+export interface BaseTool<Input = any> {
+  /**
+   * This could be undefined if the implementation is not provided,
+   *  which might be the case when communicating with a llm.
+   *
+   * @return string - the output of the tool, should be string in any case for LLM input.
+   */
+  call?: (input: Input) => string | Promise<string>;
+  metadata: // if user input any, we cannot check the schema
+  Input extends Known ? ToolMetadata<JSONSchemaType<Input>> : ToolMetadata;
 }
 
+export type BaseToolWithCall<Input = any> = Omit<BaseTool<Input>, "call"> & {
+  call: NonNullable<Pick<BaseTool<Input>, "call">["call"]>;
+};
+
 /**
  * An OutputParser is used to extract structured data from the raw output of the LLM.
  */
@@ -55,19 +88,6 @@ export interface StructuredOutput<T> {
   parsedOutput: T;
 }
 
-export type ToolParameters = {
-  type: string | "object";
-  properties: Record<string, { type: string; description?: string }>;
-  required?: string[];
-};
-
-export interface ToolMetadata {
-  description: string;
-  name: string;
-  parameters?: ToolParameters;
-  argsKwargs?: Record<string, any>;
-}
-
 export type ToolMetadataOnlyDescription = Pick<ToolMetadata, "description">;
 
 export class QueryBundle {

packages/core/tests/agent/OpenAIAgent.test.ts

@@ -1,61 +0,0 @@
-import { Settings } from "llamaindex";
-import { OpenAIAgent } from "llamaindex/agent/index";
-import { OpenAI } from "llamaindex/llm/index";
-import { FunctionTool } from "llamaindex/tools/index";
-import { beforeEach, describe, expect, it } from "vitest";
-import { mockLlmToolCallGeneration } from "../utility/mockOpenAI.js";
-
-// Define a function to sum two numbers
-function sumNumbers({ a, b }: { a: number; b: number }): number {
-  return a + b;
-}
-
-const sumJSON = {
-  type: "object",
-  properties: {
-    a: {
-      type: "number",
-      description: "The first number",
-    },
-    b: {
-      type: "number",
-      description: "The second number",
-    },
-  },
-  required: ["a", "b"],
-};
-
-describe("OpenAIAgent", () => {
-  let openaiAgent: OpenAIAgent;
-
-  beforeEach(() => {
-    const languageModel = new OpenAI({
-      model: "gpt-3.5-turbo",
-    });
-
-    Settings.llm = languageModel;
-
-    mockLlmToolCallGeneration({
-      languageModel,
-    });
-
-    const sumFunctionTool = new FunctionTool(sumNumbers, {
-      name: "sumNumbers",
-      description: "Use this function to sum two numbers",
-      parameters: sumJSON,
-    });
-
-    openaiAgent = new OpenAIAgent({
-      tools: [sumFunctionTool],
-      llm: languageModel,
-    });
-  });
-
-  it("should be able to chat with agent", async () => {
-    const response = await openaiAgent.chat({
-      message: "how much is 1 + 1?",
-    });
-
-    expect(String(response)).toEqual("The sum is 2");
-  });
-});

packages/core/tests/objects/ObjectIndex.test.ts

@@ -17,7 +17,7 @@ describe("ObjectIndex", () => {
   });
 
   test("test_object_with_tools", async () => {
-    const tool1 = new FunctionTool((x: any) => x, {
+    const tool1 = new FunctionTool(({ x }: { x: string }) => x, {
       name: "test_tool",
       description: "test tool",
       parameters: {
@@ -27,10 +27,11 @@ describe("ObjectIndex", () => {
             type: "string",
           },
         },
+        required: ["x"],
       },
     });
 
-    const tool2 = new FunctionTool((x: any) => x, {
+    const tool2 = new FunctionTool(({ x }: { x: string }) => x, {
       name: "test_tool_2",
       description: "test tool 2",
       parameters: {
@@ -40,6 +41,7 @@ describe("ObjectIndex", () => {
             type: "string",
           },
         },
+        required: ["x"],
       },
     });
 
@@ -62,7 +64,7 @@ describe("ObjectIndex", () => {
   });
 
   test("add a new object", async () => {
-    const tool1 = new FunctionTool((x: any) => x, {
+    const tool1 = new FunctionTool(({ x }: { x: string }) => x, {
       name: "test_tool",
       description: "test tool",
       parameters: {
@@ -72,10 +74,11 @@ describe("ObjectIndex", () => {
             type: "string",
           },
         },
+        required: ["x"],
       },
     });
 
-    const tool2 = new FunctionTool((x: any) => x, {
+    const tool2 = new FunctionTool(({ x }: { x: string }) => x, {
       name: "test_tool_2",
       description: "test tool 2",
       parameters: {
@@ -85,6 +88,7 @@ describe("ObjectIndex", () => {
             type: "string",
           },
         },
+        required: ["x"],
       },
     });
 

packages/core/tests/tools.test.ts

@@ -0,0 +1,32 @@
+import { ToolsFactory } from "llamaindex/tools/ToolsFactory";
+import { WikipediaTool } from "llamaindex/tools/WikipediaTool";
+import { assertType, describe, test } from "vitest";
+
+describe("ToolsFactory", async () => {
+  test("createTool", async () => {
+    await ToolsFactory.createTool(ToolsFactory.Tools.Wikipedia, {
+      metadata: {
+        name: "wikipedia_tool",
+        description: "A tool that uses a query engine to search Wikipedia.",
+      },
+    });
+  });
+  test("createTools", async () => {
+    await ToolsFactory.createTools({
+      [ToolsFactory.Tools.Wikipedia]: {
+        metadata: {
+          name: "wikipedia_tool",
+          description: "A tool that uses a query engine to search Wikipedia.",
+        },
+      },
+    });
+  });
+  test("type", () => {
+    assertType<
+      (
+        key: ToolsFactory.Tools.Wikipedia,
+        params: ConstructorParameters<typeof WikipediaTool>[0],
+      ) => Promise<WikipediaTool>
+    >(ToolsFactory.createTool<ToolsFactory.Tools.Wikipedia>);
+  });
+});

packages/core/tests/tools/Tools.test.ts

@@ -1,46 +0,0 @@
-import { FunctionTool, ToolOutput } from "llamaindex/tools/index";
-import { callToolWithErrorHandling } from "llamaindex/tools/utils";
-import { describe, expect, it } from "vitest";
-
-function sumNumbers({ a, b }: { a: number; b: number }): number {
-  return a + b;
-}
-
-const sumJSON = {
-  type: "object",
-  properties: {
-    a: {
-      type: "number",
-      description: "The first number",
-    },
-    b: {
-      type: "number",
-      description: "The second number",
-    },
-  },
-  required: ["a", "b"],
-};
-
-describe("Tools", () => {
-  it("should be able to call a tool with a common JSON", async () => {
-    const tool = new FunctionTool(sumNumbers, {
-      name: "sumNumbers",
-      description: "Use this function to sum two numbers",
-      parameters: sumJSON,
-    });
-
-    const response = await callToolWithErrorHandling(tool, {
-      a: 1,
-      b: 2,
-    });
-
-    expect(response).toEqual(
-      new ToolOutput(
-        response.content,
-        tool.metadata.name,
-        { a: 1, b: 2 },
-        response.content,
-      ),
-    );
-  });
-});

packages/core/tests/type.test.ts

@@ -0,0 +1,56 @@
+import type { ChatMessage, MessageContent, MessageType } from "llamaindex";
+import { expectTypeOf, test } from "vitest";
+import type { ChatResponse } from "../src/index.js";
+
+test("chat message type", () => {
+  // if generic is not provided, `options` is not required
+  expectTypeOf<ChatMessage>().toMatchTypeOf<{
+    content: MessageContent;
+    role: MessageType;
+  }>();
+  expectTypeOf<ChatMessage>().toMatchTypeOf<{
+    content: MessageContent;
+    role: MessageType;
+    options?: object;
+  }>();
+  expectTypeOf<ChatMessage>().not.toMatchTypeOf<{
+    content: MessageContent;
+    role: MessageType;
+    options: Record<string, unknown>;
+  }>();
+  type Options = {
+    a: string;
+    b: number;
+  };
+  expectTypeOf<ChatMessage<Options>>().toMatchTypeOf<{
+    content: MessageContent;
+    role: MessageType;
+    options?: Options;
+  }>();
+});
+
+test("chat response type", () => {
+  // if generic is not provided, `options` is not required
+  expectTypeOf<ChatResponse>().toMatchTypeOf<{
+    message: ChatMessage;
+    raw: object | null;
+  }>();
+  expectTypeOf<ChatResponse>().toMatchTypeOf<{
+    message: ChatMessage;
+    raw: object | null;
+    options?: Record<string, unknown>;
+  }>();
+  expectTypeOf<ChatResponse>().not.toMatchTypeOf<{
+    message: ChatMessage;
+    raw: object | null;
+    options: Record<string, unknown>;
+  }>();
+  type Options = {
+    a: string;
+    b: number;
+  };
+  expectTypeOf<ChatResponse<Options>>().toMatchTypeOf<{
+    message: ChatMessage<Options>;
+    raw: object | null;
+  }>();
+});