refactor: move ollama standalone (#1253)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

.github/workflows/test.yml

@@ -151,6 +151,9 @@ jobs:
       - name: Pack @llamaindex/groq
         run: pnpm pack --pack-destination ${{ runner.temp }}
         working-directory: packages/llm/groq
+      - name: Pack @llamaindex/ollama
+        run: pnpm pack --pack-destination ${{ runner.temp }}
+        working-directory: packages/llm/ollama
       - name: Pack @llamaindex/core
         run: pnpm pack --pack-destination ${{ runner.temp }}
         working-directory: packages/core

README.md

@@ -167,11 +167,13 @@ export async function chatWithAgent(
       // ... adding your tools here
     ],
   });
-  const responseStream = await agent.chat({
-    stream: true,
-    message: question,
-    chatHistory: prevMessages,
-  });
+  const responseStream = await agent.chat(
+    {
+      message: question,
+      chatHistory: prevMessages,
+    },
+    true,
+  );
   const uiStream = createStreamableUI(<div>loading...</div>);
   responseStream
     .pipeTo(

apps/docs/CHANGELOG.md

@@ -1,5 +1,41 @@
 # docs
 
+## 0.0.78
+
+### Patch Changes
+
+- llamaindex@0.6.9
+
+## 0.0.77
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+
+## 0.0.76
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+
+## 0.0.75
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+
+## 0.0.74
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+
 ## 0.0.73
 
 ### Patch Changes

apps/docs/package.json

@@ -1,6 +1,6 @@
 {
   "name": "docs",
-  "version": "0.0.73",
+  "version": "0.0.78",
   "private": true,
   "scripts": {
     "docusaurus": "docusaurus",

examples/anthropic/chat_interactive.ts

@@ -18,7 +18,7 @@ import readline from "node:readline/promises";
   });
   const chatEngine = new SimpleChatEngine({
     llm,
-    chatHistory,
+    memory: chatHistory,
   });
   const rl = readline.createInterface({ input, output });
 

examples/huggingface/embedding.ts

@@ -27,12 +27,10 @@ async function main() {
 
   // Query the index
   const queryEngine = index.asQueryEngine();
-  const stream = await queryEngine.query(
-    {
-      query: "What did the author do in college?",
-    },
-    true,
-  );
+  const stream = await queryEngine.query({
+    query: "What did the author do in college?",
+    stream: true,
+  });
 
   // Output response
   for await (const chunk of stream) {

examples/huggingface/embeddingApi.ts

@@ -37,12 +37,10 @@ async function main() {
 
   // Query the index
   const queryEngine = index.asQueryEngine();
-  const stream = await queryEngine.query(
-    {
-      query: "What did the author do in college?",
-    },
-    true,
-  );
+  const stream = await queryEngine.query({
+    query: "What did the author do in college?",
+    stream: true,
+  });
 
   // Output response
   for await (const chunk of stream) {

examples/multimodal/context.ts

@@ -1,4 +1,5 @@
 // call pnpm tsx multimodal/load.ts first to init the storage
+import { extractText } from "@llamaindex/core/utils";
 import {
   ContextChatEngine,
   NodeWithScore,
@@ -25,8 +26,9 @@ Settings.callbackManager.on("retrieve-end", (event) => {
   const textNodes = nodes.filter(
     (node: NodeWithScore) => node.node.type === ObjectType.TEXT,
   );
+  const text = extractText(query);
   console.log(
-    `Retrieved ${textNodes.length} text nodes and ${imageNodes.length} image nodes for query: ${query}`,
+    `Retrieved ${textNodes.length} text nodes and ${imageNodes.length} image nodes for query: ${text}`,
   );
 });
 

examples/multimodal/rag.ts

@@ -1,3 +1,4 @@
+import { extractText } from "@llamaindex/core/utils";
 import {
   getResponseSynthesizer,
   OpenAI,
@@ -16,7 +17,8 @@ Settings.llm = new OpenAI({ model: "gpt-4-turbo", maxTokens: 512 });
 // Update callbackManager
 Settings.callbackManager.on("retrieve-end", (event) => {
   const { nodes, query } = event.detail;
-  console.log(`Retrieved ${nodes.length} nodes for query: ${query}`);
+  const text = extractText(query);
+  console.log(`Retrieved ${nodes.length} nodes for query: ${text}`);
 });
 
 async function main() {
@@ -30,12 +32,10 @@ async function main() {
     responseSynthesizer: getResponseSynthesizer("multi_modal"),
     retriever: index.asRetriever({ topK: { TEXT: 3, IMAGE: 1 } }),
   });
-  const stream = await queryEngine.query(
-    {
-      query: "Tell me more about Vincent van Gogh's famous paintings",
-    },
-    true,
-  );
+  const stream = await queryEngine.query({
+    query: "Tell me more about Vincent van Gogh's famous paintings",
+    stream: true,
+  });
   for await (const chunk of stream) {
     process.stdout.write(chunk.response);
   }

examples/pg-vector-store/load-docs.ts

@@ -40,7 +40,11 @@ async function main(args: any) {
     const rdr = new SimpleDirectoryReader(callback);
     const docs = await rdr.loadData({ directoryPath: sourceDir });
 
-    const pgvs = new PGVectorStore();
+    const pgvs = new PGVectorStore({
+      clientConfig: {
+        connectionString: process.env.PG_CONNECTION_STRING,
+      },
+    });
     pgvs.setCollection(sourceDir);
     await pgvs.clearCollection();
 

examples/pg-vector-store/query.ts

@@ -7,7 +7,11 @@ async function main() {
   });
 
   try {
-    const pgvs = new PGVectorStore();
+    const pgvs = new PGVectorStore({
+      clientConfig: {
+        connectionString: process.env.PG_CONNECTION_STRING,
+      },
+    });
     // Optional - set your collection name, default is no filter on this field.
     // pgvs.setCollection();
 

packages/autotool/CHANGELOG.md

@@ -1,5 +1,41 @@
 # @llamaindex/autotool
 
+## 3.0.9
+
+### Patch Changes
+
+- llamaindex@0.6.9
+
+## 3.0.8
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+
+## 3.0.7
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+
+## 3.0.6
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+
+## 3.0.5
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+
 ## 3.0.4
 
 ### Patch Changes

packages/autotool/examples/01_node/CHANGELOG.md

@@ -1,5 +1,46 @@
 # @llamaindex/autotool-01-node-example
 
+## 0.0.18
+
+### Patch Changes
+
+- llamaindex@0.6.9
+- @llamaindex/autotool@3.0.9
+
+## 0.0.17
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+  - @llamaindex/autotool@3.0.8
+
+## 0.0.16
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+  - @llamaindex/autotool@3.0.7
+
+## 0.0.15
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+  - @llamaindex/autotool@3.0.6
+
+## 0.0.14
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+  - @llamaindex/autotool@3.0.5
+
 ## 0.0.13
 
 ### Patch Changes

packages/autotool/examples/01_node/package.json

@@ -13,5 +13,5 @@
   "scripts": {
     "start": "node --import tsx --import @llamaindex/autotool/node ./src/index.ts"
   },
-  "version": "0.0.13"
+  "version": "0.0.18"
 }

packages/autotool/examples/02_nextjs/CHANGELOG.md

@@ -1,5 +1,46 @@
 # @llamaindex/autotool-02-next-example
 
+## 0.1.62
+
+### Patch Changes
+
+- llamaindex@0.6.9
+- @llamaindex/autotool@3.0.9
+
+## 0.1.61
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+  - @llamaindex/autotool@3.0.8
+
+## 0.1.60
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+  - @llamaindex/autotool@3.0.7
+
+## 0.1.59
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+  - @llamaindex/autotool@3.0.6
+
+## 0.1.58
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+  - @llamaindex/autotool@3.0.5
+
 ## 0.1.57
 
 ### Patch Changes

packages/autotool/examples/02_nextjs/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/autotool-02-next-example",
   "private": true,
-  "version": "0.1.57",
+  "version": "0.1.62",
   "scripts": {
     "dev": "next dev",
     "build": "next build",

packages/autotool/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/autotool",
   "type": "module",
-  "version": "3.0.4",
+  "version": "3.0.9",
   "description": "auto transpile your JS function to LLM Agent compatible",
   "files": [
     "dist",

packages/cloud/CHANGELOG.md

@@ -1,5 +1,11 @@
 # @llamaindex/cloud
 
+## 0.2.8
+
+### Patch Changes
+
+- ac41ed3: feat: bump cloud sdk version
+
 ## 0.2.7
 
 ### Patch Changes

packages/cloud/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/cloud",
-  "version": "0.2.7",
+  "version": "0.2.8",
   "type": "module",
   "license": "MIT",
   "scripts": {

packages/cloud/src/reader.ts

@@ -170,6 +170,15 @@ export class LlamaParseReader extends FileReader {
   vendorMultimodalModelName?: string | undefined;
   // The API key for the multimodal API. Can also be set as an env variable: LLAMA_CLOUD_VENDOR_MULTIMODAL_API_KEY
   vendorMultimodalApiKey?: string | undefined;
+
+  webhookUrl?: string | undefined;
+  premiumMode?: boolean | undefined;
+  takeScreenshot?: boolean | undefined;
+  disableOcr?: boolean | undefined;
+  disableReconstruction?: boolean | undefined;
+  inputS3Path?: string | undefined;
+  outputS3PathPrefix?: string | undefined;
+
   // numWorkers is implemented in SimpleDirectoryReader
   stdout?: WriteStream | undefined;
 
@@ -258,13 +267,13 @@ export class LlamaParseReader extends FileReader {
       use_vendor_multimodal_model: this.useVendorMultimodalModel,
       vendor_multimodal_model_name: this.vendorMultimodalModelName,
       vendor_multimodal_api_key: this.vendorMultimodalApiKey,
-      // fixme: does these fields need to be set?
-      webhook_url: undefined,
-      take_screenshot: undefined,
-      disable_ocr: undefined,
-      disable_reconstruction: undefined,
-      input_s3_path: undefined,
-      output_s3_path_prefix: undefined,
+      premium_mode: this.premiumMode,
+      webhook_url: this.webhookUrl,
+      take_screenshot: this.takeScreenshot,
+      disable_ocr: this.disableOcr,
+      disable_reconstruction: this.disableReconstruction,
+      input_s3_path: this.inputS3Path,
+      output_s3_path_prefix: this.outputS3PathPrefix,
     } satisfies {
       [Key in keyof Body_upload_file_api_v1_parsing_upload_post]-?:
         | Body_upload_file_api_v1_parsing_upload_post[Key]

packages/cloud/turbo.json

@@ -0,0 +1,8 @@
+{
+  "extends": ["//"],
+  "tasks": {
+    "build": {
+      "outputs": ["dist/**", "src/client/**"]
+    }
+  }
+}

packages/community/CHANGELOG.md

@@ -1,5 +1,20 @@
 # @llamaindex/community
 
+## 0.0.40
+
+### Patch Changes
+
+- 50e6b57: feat: add Amazon Bedrock Retriever
+- Updated dependencies [8b7fdba]
+  - @llamaindex/core@0.2.6
+
+## 0.0.39
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+  - @llamaindex/core@0.2.5
+
 ## 0.0.38
 
 ### Patch Changes

packages/community/README.md

@@ -7,6 +7,7 @@
 - Bedrock support for the Anthropic Claude Models [usage](https://ts.llamaindex.ai/modules/llms/available_llms/bedrock)
 - Bedrock support for the Meta LLama 2, 3 and 3.1 Models [usage](https://ts.llamaindex.ai/modules/llms/available_llms/bedrock)
 - Meta LLama3.1 405b tool call support
+- Bedrock support for querying Knowledge Base
 
 ## LICENSE
 

packages/community/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/community",
   "description": "Community package for LlamaIndexTS",
-  "version": "0.0.38",
+  "version": "0.0.40",
   "type": "module",
   "types": "dist/type/index.d.ts",
   "main": "dist/cjs/index.js",
@@ -47,6 +47,7 @@
   },
   "dependencies": {
     "@aws-sdk/client-bedrock-runtime": "^3.642.0",
+    "@aws-sdk/client-bedrock-agent-runtime": "^3.642.0",
     "@llamaindex/core": "workspace:*",
     "@llamaindex/env": "workspace:*"
   }

packages/community/src/index.ts

@@ -3,3 +3,4 @@ export {
   BEDROCK_MODEL_MAX_TOKENS,
   Bedrock,
 } from "./llm/bedrock/index.js";
+export { AmazonKnowledgeBaseRetriever } from "./retrievers/bedrock.js";

packages/community/src/retrievers/bedrock.ts

@@ -0,0 +1,165 @@
+import type { KnowledgeBaseVectorSearchConfiguration } from "@aws-sdk/client-bedrock-agent-runtime";
+import {
+  BedrockAgentRuntimeClient,
+  type BedrockAgentRuntimeClientConfig,
+  type RetrievalFilter,
+  RetrieveCommand,
+  type SearchType,
+} from "@aws-sdk/client-bedrock-agent-runtime";
+import type { QueryBundle } from "@llamaindex/core/query-engine";
+import { BaseRetriever } from "@llamaindex/core/retriever";
+import { Document, type NodeWithScore } from "@llamaindex/core/schema";
+import { extractText } from "@llamaindex/core/utils";
+
+/**
+ * Interface for the arguments required to initialize an
+ * AmazonKnowledgeBaseRetriever instance.
+ */
+export interface AmazonKnowledgeBaseRetrieverArgs {
+  knowledgeBaseId: string;
+  topK: number;
+  region: string;
+  clientOptions?: BedrockAgentRuntimeClientConfig;
+  filter?: RetrievalFilter;
+  overrideSearchType?: SearchType;
+}
+
+/**
+ * Class for interacting with Amazon Bedrock Knowledge Bases, a RAG workflow oriented service
+ * Extends the BaseRetriever class.
+ * @example
+ * ```typescript
+ * const retriever = new AmazonKnowledgeBaseRetriever({
+ *   topK: 10,
+ *   knowledgeBaseId: "YOUR_KNOWLEDGE_BASE_ID",
+ *   region: "us-east-2",
+ *   clientOptions: {
+ *     credentials: {
+ *       accessKeyId: "YOUR_ACCESS_KEY_ID",
+ *       secretAccessKey: "YOUR_SECRET_ACCESS_KEY",
+ *     },
+ *   },
+ * });
+ *
+ * const docs = await retriever.retrieve({query: "How are clouds formed?"});
+ * ```
+ */
+export class AmazonKnowledgeBaseRetriever extends BaseRetriever {
+  static lc_name() {
+    return "AmazonKnowledgeBaseRetriever";
+  }
+
+  lc_namespace = ["llamaindex", "retrievers", "amazon_bedrock_knowledge_base"];
+
+  knowledgeBaseId: string;
+
+  topK: number;
+
+  bedrockAgentRuntimeClient: BedrockAgentRuntimeClient;
+
+  filter: RetrievalFilter | undefined;
+
+  overrideSearchType: SearchType | undefined;
+
+  constructor({
+    knowledgeBaseId,
+    topK = 10,
+    clientOptions,
+    region,
+    filter,
+    overrideSearchType,
+  }: AmazonKnowledgeBaseRetrieverArgs) {
+    super();
+
+    this.topK = topK;
+    this.filter = filter;
+    this.overrideSearchType = overrideSearchType;
+    this.bedrockAgentRuntimeClient = new BedrockAgentRuntimeClient({
+      region,
+      ...clientOptions,
+    });
+    this.knowledgeBaseId = knowledgeBaseId;
+  }
+
+  /**
+   * Cleans the result text by replacing sequences of whitespace with a
+   * single space and removing ellipses.
+   * @param resText The result text to clean.
+   * @returns The cleaned result text.
+   */
+  cleanResult(resText: string) {
+    const res = resText.replace(/\s+/g, " ").replace(/\.\.\./g, "");
+    return res;
+  }
+
+  async queryKnowledgeBase(
+    query: QueryBundle,
+    topK: number,
+    filter?: RetrievalFilter,
+    overrideSearchType?: SearchType,
+  ): Promise<NodeWithScore[]> {
+    const retrieveCommand = new RetrieveCommand({
+      knowledgeBaseId: this.knowledgeBaseId,
+      retrievalQuery: {
+        text: extractText(query),
+      },
+      retrievalConfiguration: {
+        vectorSearchConfiguration: {
+          numberOfResults: topK,
+          overrideSearchType,
+          filter,
+        } as KnowledgeBaseVectorSearchConfiguration,
+      },
+    });
+
+    const retrieveResponse =
+      await this.bedrockAgentRuntimeClient.send(retrieveCommand);
+
+    return (
+      retrieveResponse.retrievalResults?.map((result) => {
+        let source;
+        switch (result.location?.type) {
+          case "CONFLUENCE":
+            source = result.location?.confluenceLocation?.url;
+            break;
+          case "S3":
+            source = result.location?.s3Location?.uri;
+            break;
+          case "SALESFORCE":
+            source = result.location?.salesforceLocation?.url;
+            break;
+          case "SHAREPOINT":
+            source = result.location?.sharePointLocation?.url;
+            break;
+          case "WEB":
+            source = result.location?.webLocation?.url;
+            break;
+          default:
+            source = result.location?.s3Location?.uri;
+            break;
+        }
+
+        return {
+          node: new Document({
+            text: this.cleanResult(result.content?.text || ""),
+            metadata: {
+              source,
+              score: result.score,
+              ...result.metadata,
+            },
+          }),
+          score: result.score ?? 1.0,
+        };
+      }) ?? []
+    );
+  }
+
+  async _retrieve(query: QueryBundle): Promise<NodeWithScore[]> {
+    return await this.queryKnowledgeBase(
+      query,
+      this.topK,
+      this.filter,
+      this.overrideSearchType,
+    );
+  }
+}

packages/core/CHANGELOG.md

@@ -1,5 +1,20 @@
 # @llamaindex/core
 
+## 0.2.6
+
+### Patch Changes
+
+- 8b7fdba: refactor: move chat engine & retriever into core.
+
+  - `chatHistory` in BaseChatEngine now returns `ChatMessage[] | Promise<ChatMessage[]>`, instead of `BaseMemory`
+  - update `retrieve-end` type
+
+## 0.2.5
+
+### Patch Changes
+
+- d902cc3: Fix context not being sent using ContextChatEngine
+
 ## 0.2.4
 
 ### Patch Changes

packages/core/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/core",
   "type": "module",
-  "version": "0.2.4",
+  "version": "0.2.6",
   "description": "LlamaIndex Core Module",
   "exports": {
     "./node-parser": {
@@ -199,6 +199,34 @@
         "types": "./dist/response-synthesizers/index.d.ts",
         "default": "./dist/response-synthesizers/index.js"
       }
+    },
+    "./chat-engine": {
+      "require": {
+        "types": "./dist/chat-engine/index.d.cts",
+        "default": "./dist/chat-engine/index.cjs"
+      },
+      "import": {
+        "types": "./dist/chat-engine/index.d.ts",
+        "default": "./dist/chat-engine/index.js"
+      },
+      "default": {
+        "types": "./dist/chat-engine/index.d.ts",
+        "default": "./dist/chat-engine/index.js"
+      }
+    },
+    "./retriever": {
+      "require": {
+        "types": "./dist/retriever/index.d.cts",
+        "default": "./dist/retriever/index.cjs"
+      },
+      "import": {
+        "types": "./dist/retriever/index.d.ts",
+        "default": "./dist/retriever/index.js"
+      },
+      "default": {
+        "types": "./dist/retriever/index.d.ts",
+        "default": "./dist/retriever/index.js"
+      }
     }
   },
   "files": [

packages/core/src/chat-engine/index.ts

@@ -0,0 +1,36 @@
+import type { ChatMessage, MessageContent } from "../llms";
+import type { BaseMemory } from "../memory";
+import { EngineResponse } from "../schema";
+
+export interface BaseChatEngineParams<
+  AdditionalMessageOptions extends object = object,
+> {
+  message: MessageContent;
+  /**
+   * Optional chat history if you want to customize the chat history.
+   */
+  chatHistory?:
+    | ChatMessage<AdditionalMessageOptions>[]
+    | BaseMemory<AdditionalMessageOptions>;
+}
+
+export interface StreamingChatEngineParams<
+  AdditionalMessageOptions extends object = object,
+> extends BaseChatEngineParams<AdditionalMessageOptions> {
+  stream: true;
+}
+
+export interface NonStreamingChatEngineParams<
+  AdditionalMessageOptions extends object = object,
+> extends BaseChatEngineParams<AdditionalMessageOptions> {
+  stream?: false;
+}
+
+export abstract class BaseChatEngine {
+  abstract chat(params: NonStreamingChatEngineParams): Promise<EngineResponse>;
+  abstract chat(
+    params: StreamingChatEngineParams,
+  ): Promise<AsyncIterable<EngineResponse>>;
+
+  abstract chatHistory: ChatMessage[] | Promise<ChatMessage[]>;
+}

packages/core/src/global/settings/callback-manager.ts

@@ -11,6 +11,7 @@ import type {
   SynthesizeEndEvent,
   SynthesizeStartEvent,
 } from "../../response-synthesizers";
+import type { RetrieveEndEvent, RetrieveStartEvent } from "../../retriever";
 import { TextNode } from "../../schema";
 import { EventCaller, getEventCaller } from "../../utils";
 import type { UUID } from "../type";
@@ -69,6 +70,8 @@ export interface LlamaIndexEventMaps {
   "query-end": QueryEndEvent;
   "synthesize-start": SynthesizeStartEvent;
   "synthesize-end": SynthesizeEndEvent;
+  "retrieve-start": RetrieveStartEvent;
+  "retrieve-end": RetrieveEndEvent;
 }
 
 export class LlamaIndexCustomEvent<T = any> extends CustomEvent<T> {

packages/core/src/memory/base.ts

@@ -1,5 +1,5 @@
 import { Settings } from "../global";
-import type { ChatMessage, MessageContent } from "../llms";
+import type { ChatMessage } from "../llms";
 import { type BaseChatStore, SimpleChatStore } from "../storage/chat-store";
 import { extractText } from "../utils";
 
@@ -12,15 +12,36 @@ export const DEFAULT_CHAT_STORE_KEY = "chat_history";
 export abstract class BaseMemory<
   AdditionalMessageOptions extends object = object,
 > {
+  /**
+   * Retrieves messages from the memory, optionally including transient messages.
+   * Compared to getAllMessages, this method a) allows for transient messages to be included in the retrieval and b) may return a subset of the total messages by applying a token limit.
+   * @param transientMessages Optional array of temporary messages to be included in the retrieval.
+   * These messages are not stored in the memory but are considered for the current interaction.
+   * @returns An array of chat messages, either synchronously or as a Promise.
+   */
   abstract getMessages(
-    input?: MessageContent | undefined,
+    transientMessages?: ChatMessage<AdditionalMessageOptions>[] | undefined,
   ):
     | ChatMessage<AdditionalMessageOptions>[]
     | Promise<ChatMessage<AdditionalMessageOptions>[]>;
+
+  /**
+   * Retrieves all messages stored in the memory.
+   * @returns An array of all chat messages, either synchronously or as a Promise.
+   */
   abstract getAllMessages():
     | ChatMessage<AdditionalMessageOptions>[]
     | Promise<ChatMessage<AdditionalMessageOptions>[]>;
+
+  /**
+   * Adds a new message to the memory.
+   * @param messages The chat message to be added to the memory.
+   */
   abstract put(messages: ChatMessage<AdditionalMessageOptions>): void;
+
+  /**
+   * Clears all messages from the memory.
+   */
   abstract reset(): void;
 
   protected _tokenCountForMessages(messages: ChatMessage[]): number {

packages/core/src/memory/chat-memory-buffer.ts

@@ -1,5 +1,5 @@
 import { Settings } from "../global";
-import type { ChatMessage, LLM, MessageContent } from "../llms";
+import type { ChatMessage, LLM } from "../llms";
 import { type BaseChatStore } from "../storage/chat-store";
 import { BaseChatStoreMemory, DEFAULT_TOKEN_LIMIT_RATIO } from "./base";
 
@@ -34,7 +34,7 @@ export class ChatMemoryBuffer<
   }
 
   getMessages(
-    input?: MessageContent | undefined,
+    transientMessages?: ChatMessage<AdditionalMessageOptions>[] | undefined,
     initialTokenCount: number = 0,
   ) {
     const messages = this.getAllMessages();
@@ -43,16 +43,22 @@ export class ChatMemoryBuffer<
       throw new Error("Initial token count exceeds token limit");
     }
 
-    let messageCount = messages.length;
-    let currentMessages = messages.slice(-messageCount);
-    let tokenCount = this._tokenCountForMessages(messages) + initialTokenCount;
+    // Add input messages as transient messages
+    const messagesWithInput = transientMessages
+      ? [...transientMessages, ...messages]
+      : messages;
+
+    let messageCount = messagesWithInput.length;
+    let currentMessages = messagesWithInput.slice(-messageCount);
+    let tokenCount =
+      this._tokenCountForMessages(messagesWithInput) + initialTokenCount;
 
     while (tokenCount > this.tokenLimit && messageCount > 1) {
       messageCount -= 1;
-      if (messages.at(-messageCount)!.role === "assistant") {
+      if (messagesWithInput.at(-messageCount)!.role === "assistant") {
         messageCount -= 1;
       }
-      currentMessages = messages.slice(-messageCount);
+      currentMessages = messagesWithInput.slice(-messageCount);
       tokenCount =
         this._tokenCountForMessages(currentMessages) + initialTokenCount;
     }
@@ -60,6 +66,6 @@ export class ChatMemoryBuffer<
     if (tokenCount > this.tokenLimit && messageCount <= 0) {
       return [];
     }
-    return messages.slice(-messageCount);
+    return messagesWithInput.slice(-messageCount);
   }
 }

packages/core/src/memory/summary-memory.ts

@@ -114,18 +114,22 @@ export class ChatSummaryMemoryBuffer extends BaseMemory {
     }
   }
 
-  private calcCurrentRequestMessages() {
-    // TODO: check order: currently, we're sending:
+  private calcCurrentRequestMessages(transientMessages?: ChatMessage[]) {
+    // currently, we're sending:
     // system messages first, then transient messages and then the messages that describe the conversation so far
-    return [...this.systemMessages, ...this.calcConversationMessages(true)];
+    return [
+      ...this.systemMessages,
+      ...(transientMessages ? transientMessages : []),
+      ...this.calcConversationMessages(true),
+    ];
   }
 
   reset() {
     this.messages = [];
   }
 
-  async getMessages(): Promise<ChatMessage[]> {
-    const requestMessages = this.calcCurrentRequestMessages();
+  async getMessages(transientMessages?: ChatMessage[]): Promise<ChatMessage[]> {
+    const requestMessages = this.calcCurrentRequestMessages(transientMessages);
 
     // get tokens of current request messages and the transient messages
     const tokens = requestMessages.reduce(
@@ -149,7 +153,7 @@ export class ChatSummaryMemoryBuffer extends BaseMemory {
       // TODO: we still might have too many tokens
       // e.g. too large system messages or transient messages
       // how should we deal with that?
-      return this.calcCurrentRequestMessages();
+      return this.calcCurrentRequestMessages(transientMessages);
     }
     return requestMessages;
   }

packages/core/src/query-engine/base.ts

@@ -2,7 +2,7 @@ import { randomUUID } from "@llamaindex/env";
 import { Settings } from "../global";
 import type { MessageContent } from "../llms";
 import { PromptMixin } from "../prompts";
-import { EngineResponse } from "../schema";
+import { EngineResponse, type NodeWithScore } from "../schema";
 import { wrapEventCaller } from "../utils";
 
 /**
@@ -18,6 +18,18 @@ export type QueryBundle = {
 
 export type QueryType = string | QueryBundle;
 
+export type BaseQueryParams = {
+  query: QueryType;
+};
+
+export interface StreamingQueryParams extends BaseQueryParams {
+  stream: true;
+}
+
+export interface NonStreamingQueryParams extends BaseQueryParams {
+  stream?: false;
+}
+
 export type QueryFn = (
   strOrQueryBundle: QueryType,
   stream?: boolean,
@@ -28,23 +40,26 @@ export abstract class BaseQueryEngine extends PromptMixin {
     super();
   }
 
-  query(
-    strOrQueryBundle: QueryType,
-    stream: true,
-  ): Promise<AsyncIterable<EngineResponse>>;
-  query(strOrQueryBundle: QueryType, stream?: false): Promise<EngineResponse>;
+  async retrieve(params: QueryType): Promise<NodeWithScore[]> {
+    throw new Error(
+      "This query engine does not support retrieve, use query directly",
+    );
+  }
+
+  query(params: StreamingQueryParams): Promise<AsyncIterable<EngineResponse>>;
+  query(params: NonStreamingQueryParams): Promise<EngineResponse>;
   @wrapEventCaller
   async query(
-    strOrQueryBundle: QueryType,
-    stream = false,
+    params: StreamingQueryParams | NonStreamingQueryParams,
   ): Promise<EngineResponse | AsyncIterable<EngineResponse>> {
+    const { stream, query } = params;
     const id = randomUUID();
     const callbackManager = Settings.callbackManager;
     callbackManager.dispatchEvent("query-start", {
       id,
-      query: strOrQueryBundle,
+      query,
     });
-    const response = await this._query(strOrQueryBundle, stream);
+    const response = await this._query(query, stream);
     callbackManager.dispatchEvent("query-end", {
       id,
       response,

packages/core/src/retriever/index.ts

@@ -0,0 +1,112 @@
+import { randomUUID } from "@llamaindex/env";
+import { Settings } from "../global";
+import type { MessageContent } from "../llms";
+import { PromptMixin } from "../prompts";
+import type { QueryBundle, QueryType } from "../query-engine";
+import { BaseNode, IndexNode, type NodeWithScore, ObjectType } from "../schema";
+
+export type RetrieveParams = {
+  query: MessageContent;
+  preFilters?: unknown;
+};
+
+export type RetrieveStartEvent = {
+  id: string;
+  query: QueryBundle;
+};
+
+export type RetrieveEndEvent = {
+  id: string;
+  query: QueryBundle;
+  nodes: NodeWithScore[];
+};
+
+export abstract class BaseRetriever extends PromptMixin {
+  objectMap: Map<string, unknown> = new Map();
+
+  protected _updatePrompts() {}
+  protected _getPrompts() {
+    return {};
+  }
+
+  protected _getPromptModules() {
+    return {};
+  }
+
+  protected constructor() {
+    super();
+  }
+
+  public async retrieve(params: QueryType): Promise<NodeWithScore[]> {
+    const cb = Settings.callbackManager;
+    const queryBundle = typeof params === "string" ? { query: params } : params;
+    const id = randomUUID();
+    cb.dispatchEvent("retrieve-start", { id, query: queryBundle });
+    let response = await this._retrieve(queryBundle);
+    response = await this._handleRecursiveRetrieval(queryBundle, response);
+    cb.dispatchEvent("retrieve-end", {
+      id,
+      query: queryBundle,
+      nodes: response,
+    });
+    return response;
+  }
+
+  abstract _retrieve(params: QueryBundle): Promise<NodeWithScore[]>;
+
+  async _handleRecursiveRetrieval(
+    params: QueryBundle,
+    nodes: NodeWithScore[],
+  ): Promise<NodeWithScore[]> {
+    const retrievedNodes = [];
+    for (const { node, score = 1.0 } of nodes) {
+      if (node.type === ObjectType.INDEX) {
+        const indexNode = node as IndexNode;
+        const object = this.objectMap.get(indexNode.indexId);
+        if (object !== undefined) {
+          retrievedNodes.push(
+            ...this._retrieveFromObject(object, params, score),
+          );
+        } else {
+          retrievedNodes.push({ node, score });
+        }
+      } else {
+        retrievedNodes.push({ node, score });
+      }
+    }
+    return nodes;
+  }
+
+  _retrieveFromObject(
+    object: unknown,
+    queryBundle: QueryBundle,
+    score: number,
+  ): NodeWithScore[] {
+    if (object == null) {
+      throw new TypeError("Object is not retrievable");
+    }
+    if (typeof object !== "object") {
+      throw new TypeError("Object is not retrievable");
+    }
+    if ("node" in object && object.node instanceof BaseNode) {
+      return [
+        {
+          node: object.node,
+          score:
+            "score" in object && typeof object.score === "number"
+              ? object.score
+              : score,
+        },
+      ];
+    }
+    if (object instanceof BaseNode) {
+      return [{ node: object, score }];
+    } else {
+      // todo: support other types
+      // BaseQueryEngine
+      // BaseRetriever
+      // QueryComponent
+      throw new TypeError("Object is not retrievable");
+    }
+  }
+}

packages/core/tests/memory/chat-memory-buffer.test.ts

@@ -0,0 +1,74 @@
+import { Settings } from "@llamaindex/core/global";
+import type { ChatMessage } from "@llamaindex/core/llms";
+import { ChatMemoryBuffer } from "@llamaindex/core/memory";
+import { beforeEach, describe, expect, test } from "vitest";
+
+describe("ChatMemoryBuffer", () => {
+  beforeEach(() => {
+    // Mock the Settings.llm
+    (Settings.llm as any) = {
+      metadata: {
+        contextWindow: 1000,
+      },
+    };
+  });
+
+  test("constructor initializes with custom token limit", () => {
+    const buffer = new ChatMemoryBuffer({ tokenLimit: 500 });
+    expect(buffer.tokenLimit).toBe(500);
+  });
+
+  test("getMessages returns all messages when under token limit", () => {
+    const messages: ChatMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there!" },
+      { role: "user", content: "How are you?" },
+    ];
+    const buffer = new ChatMemoryBuffer({
+      tokenLimit: 1000,
+      chatHistory: messages,
+    });
+
+    const result = buffer.getMessages();
+    expect(result).toEqual(messages);
+  });
+
+  test("getMessages truncates messages when over token limit", () => {
+    const messages: ChatMessage[] = [
+      { role: "user", content: "This is a long message" },
+      { role: "assistant", content: "This is also a long reply" },
+      { role: "user", content: "Short" },
+    ];
+    const buffer = new ChatMemoryBuffer({
+      tokenLimit: 5, // limit to only allow the last message
+      chatHistory: messages,
+    });
+
+    const result = buffer.getMessages();
+    expect(result).toEqual([{ role: "user", content: "Short" }]);
+  });
+
+  test("getMessages handles input messages", () => {
+    const storedMessages: ChatMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there!" },
+    ];
+    const buffer = new ChatMemoryBuffer({
+      tokenLimit: 50,
+      chatHistory: storedMessages,
+    });
+
+    const inputMessages: ChatMessage[] = [
+      { role: "user", content: "New message" },
+    ];
+    const result = buffer.getMessages(inputMessages);
+    expect(result).toEqual([...inputMessages, ...storedMessages]);
+  });
+
+  test("getMessages throws error when initial token count exceeds limit", () => {
+    const buffer = new ChatMemoryBuffer({ tokenLimit: 10 });
+    expect(() => buffer.getMessages(undefined, 20)).toThrow(
+      "Initial token count exceeds token limit",
+    );
+  });
+});

packages/experimental/CHANGELOG.md

@@ -1,5 +1,41 @@
 # @llamaindex/experimental
 
+## 0.0.87
+
+### Patch Changes
+
+- llamaindex@0.6.9
+
+## 0.0.86
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+
+## 0.0.85
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+
+## 0.0.84
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+
+## 0.0.83
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+
 ## 0.0.82
 
 ### Patch Changes

packages/experimental/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/experimental",
   "description": "Experimental package for LlamaIndexTS",
-  "version": "0.0.82",
+  "version": "0.0.87",
   "type": "module",
   "types": "dist/type/index.d.ts",
   "main": "dist/cjs/index.js",

packages/llamaindex/CHANGELOG.md

@@ -1,5 +1,56 @@
 # llamaindex
 
+## 0.6.9
+
+### Patch Changes
+
+- Updated dependencies [ac41ed3]
+  - @llamaindex/cloud@0.2.8
+
+## 0.6.8
+
+### Patch Changes
+
+- 8b7fdba: refactor: move chat engine & retriever into core.
+
+  - `chatHistory` in BaseChatEngine now returns `ChatMessage[] | Promise<ChatMessage[]>`, instead of `BaseMemory`
+  - update `retrieve-end` type
+
+- Updated dependencies [8b7fdba]
+  - @llamaindex/core@0.2.6
+  - @llamaindex/openai@0.1.8
+  - @llamaindex/groq@0.0.7
+
+## 0.6.7
+
+### Patch Changes
+
+- 23bcc37: fix: add `serializer` in doc store
+
+  `PostgresDocumentStore` now will not use JSON.stringify for better performance
+
+## 0.6.6
+
+### Patch Changes
+
+- d902cc3: Fix context not being sent using ContextChatEngine
+- 025ffe6: fix: update `PostgresKVStore` constructor params
+- a659574: Adds upstash vector store as a storage
+- Updated dependencies [d902cc3]
+  - @llamaindex/core@0.2.5
+  - @llamaindex/openai@0.1.7
+  - @llamaindex/groq@0.0.6
+
+## 0.6.5
+
+### Patch Changes
+
+- e9714db: feat: update `PGVectorStore`
+
+  - move constructor parameter `config.user` | `config.database` | `config.password` | `config.connectionString` into `config.clientConfig`
+  - if you pass `pg.Client` or `pg.Pool` instance to `PGVectorStore`, move it to `config.client`, setting `config.shouldConnect` to false if it's already connected
+  - default value of `PGVectorStore.collection` is now `"data"` instead of `""` (empty string)
+
 ## 0.6.4
 
 ### Patch Changes

packages/llamaindex/e2e/examples/cloudflare-worker-agent/CHANGELOG.md

@@ -1,5 +1,41 @@
 # @llamaindex/cloudflare-worker-agent-test
 
+## 0.0.71
+
+### Patch Changes
+
+- llamaindex@0.6.9
+
+## 0.0.70
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+
+## 0.0.69
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+
+## 0.0.68
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+
+## 0.0.67
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+
 ## 0.0.66
 
 ### Patch Changes

packages/llamaindex/e2e/examples/cloudflare-worker-agent/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/cloudflare-worker-agent-test",
-  "version": "0.0.66",
+  "version": "0.0.71",
   "type": "module",
   "private": true,
   "scripts": {

packages/llamaindex/e2e/examples/llama-parse-browser/CHANGELOG.md

@@ -1,5 +1,12 @@
 # @llamaindex/llama-parse-browser-test
 
+## 0.0.4
+
+### Patch Changes
+
+- Updated dependencies [ac41ed3]
+  - @llamaindex/cloud@0.2.8
+
 ## 0.0.3
 
 ### Patch Changes

packages/llamaindex/e2e/examples/llama-parse-browser/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/llama-parse-browser-test",
   "private": true,
-  "version": "0.0.3",
+  "version": "0.0.4",
   "type": "module",
   "scripts": {
     "dev": "vite",

packages/llamaindex/e2e/examples/nextjs-agent/CHANGELOG.md

@@ -1,5 +1,41 @@
 # @llamaindex/next-agent-test
 
+## 0.1.71
+
+### Patch Changes
+
+- llamaindex@0.6.9
+
+## 0.1.70
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+
+## 0.1.69
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+
+## 0.1.68
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+
+## 0.1.67
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+
 ## 0.1.66
 
 ### Patch Changes

packages/llamaindex/e2e/examples/nextjs-agent/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/next-agent-test",
-  "version": "0.1.66",
+  "version": "0.1.71",
   "private": true,
   "scripts": {
     "dev": "next dev",

packages/llamaindex/e2e/examples/nextjs-edge-runtime/CHANGELOG.md

@@ -1,5 +1,41 @@
 # test-edge-runtime
 
+## 0.1.70
+
+### Patch Changes
+
+- llamaindex@0.6.9
+
+## 0.1.69
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+
+## 0.1.68
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+
+## 0.1.67
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+
+## 0.1.66
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+
 ## 0.1.65
 
 ### Patch Changes

packages/llamaindex/e2e/examples/nextjs-edge-runtime/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/nextjs-edge-runtime-test",
-  "version": "0.1.65",
+  "version": "0.1.70",
   "private": true,
   "scripts": {
     "dev": "next dev",

packages/llamaindex/e2e/examples/nextjs-node-runtime/CHANGELOG.md

@@ -1,5 +1,41 @@
 # @llamaindex/next-node-runtime
 
+## 0.0.52
+
+### Patch Changes
+
+- llamaindex@0.6.9
+
+## 0.0.51
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+
+## 0.0.50
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+
+## 0.0.49
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+
+## 0.0.48
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+
 ## 0.0.47
 
 ### Patch Changes

packages/llamaindex/e2e/examples/nextjs-node-runtime/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/next-node-runtime-test",
-  "version": "0.0.47",
+  "version": "0.0.52",
   "private": true,
   "scripts": {
     "dev": "next dev",

packages/llamaindex/e2e/examples/waku-query-engine/CHANGELOG.md

@@ -1,5 +1,41 @@
 # @llamaindex/waku-query-engine-test
 
+## 0.0.71
+
+### Patch Changes
+
+- llamaindex@0.6.9
+
+## 0.0.70
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - llamaindex@0.6.8
+
+## 0.0.69
+
+### Patch Changes
+
+- Updated dependencies [23bcc37]
+  - llamaindex@0.6.7
+
+## 0.0.68
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+- Updated dependencies [025ffe6]
+- Updated dependencies [a659574]
+  - llamaindex@0.6.6
+
+## 0.0.67
+
+### Patch Changes
+
+- Updated dependencies [e9714db]
+  - llamaindex@0.6.5
+
 ## 0.0.66
 
 ### Patch Changes

packages/llamaindex/e2e/examples/waku-query-engine/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/waku-query-engine-test",
-  "version": "0.0.66",
+  "version": "0.0.71",
   "type": "module",
   "private": true,
   "scripts": {

packages/llamaindex/e2e/node/vector-store/pg-vector-store.e2e.ts

@@ -9,43 +9,54 @@ import { registerTypes } from "pgvector/pg";
 
 config({ path: [".env.local", ".env", ".env.ci"] });
 
-let pgClient: pg.Client | pg.Pool;
-test.afterEach(async () => {
-  await pgClient.end();
-});
-
 const pgConfig = {
   user: process.env.POSTGRES_USER ?? "user",
   password: process.env.POSTGRES_PASSWORD ?? "password",
   database: "llamaindex_node_test",
 };
 
-await test("init with client", async () => {
-  pgClient = new pg.Client(pgConfig);
+await test("init with client", async (t) => {
+  const pgClient = new pg.Client(pgConfig);
   await pgClient.connect();
   await pgClient.query("CREATE EXTENSION IF NOT EXISTS vector");
   await registerTypes(pgClient);
-  const vectorStore = new PGVectorStore(pgClient);
+  t.after(async () => {
+    await pgClient.end();
+  });
+  const vectorStore = new PGVectorStore({
+    client: pgClient,
+    shouldConnect: false,
+  });
   assert.deepStrictEqual(await vectorStore.client(), pgClient);
 });
 
-await test("init with pool", async () => {
-  pgClient = new pg.Pool(pgConfig);
+await test("init with pool", async (t) => {
+  const pgClient = new pg.Pool(pgConfig);
   await pgClient.query("CREATE EXTENSION IF NOT EXISTS vector");
   const client = await pgClient.connect();
+  await client.query("CREATE EXTENSION IF NOT EXISTS vector");
   await registerTypes(client);
-  const vectorStore = new PGVectorStore(client);
+  t.after(async () => {
+    client.release();
+    await pgClient.end();
+  });
+  const vectorStore = new PGVectorStore({
+    shouldConnect: false,
+    client,
+  });
   assert.deepStrictEqual(await vectorStore.client(), client);
-  client.release();
 });
 
-await test("init without client", async () => {
-  const vectorStore = new PGVectorStore(pgConfig);
-  pgClient = (await vectorStore.client()) as pg.Client;
+await test("init without client", async (t) => {
+  const vectorStore = new PGVectorStore({ clientConfig: pgConfig });
+  const pgClient = (await vectorStore.client()) as pg.Client;
+  t.after(async () => {
+    await pgClient.end();
+  });
   assert.notDeepStrictEqual(pgClient, undefined);
 });
 
-await test("simple node", async () => {
+await test("simple node", async (t) => {
   const dimensions = 3;
   const schemaName =
     "llamaindex_vector_store_test_" + Math.random().toString(36).substring(7);
@@ -56,10 +67,14 @@ await test("simple node", async () => {
     embedding: [0.1, 0.2, 0.3],
   });
   const vectorStore = new PGVectorStore({
-    ...pgConfig,
+    clientConfig: pgConfig,
     dimensions,
     schemaName,
   });
+  const pgClient = (await vectorStore.client()) as pg.Client;
+  t.after(async () => {
+    await pgClient.end();
+  });
 
   await vectorStore.add([node]);
 
@@ -89,6 +104,4 @@ await test("simple node", async () => {
     });
     assert.deepStrictEqual(result.nodes, []);
   }
-
-  pgClient = (await vectorStore.client()) as pg.Client;
 });

packages/llamaindex/package.json

@@ -1,6 +1,6 @@
 {
   "name": "llamaindex",
-  "version": "0.6.4",
+  "version": "0.6.9",
   "license": "MIT",
   "type": "module",
   "keywords": [
@@ -34,6 +34,7 @@
     "@llamaindex/core": "workspace:*",
     "@llamaindex/env": "workspace:*",
     "@llamaindex/groq": "workspace:*",
+    "@llamaindex/ollama": "workspace:*",
     "@llamaindex/openai": "workspace:*",
     "@mistralai/mistralai": "^1.0.4",
     "@mixedbread-ai/sdk": "^2.2.11",
@@ -43,6 +44,7 @@
     "@types/node": "^22.5.1",
     "@types/papaparse": "^5.3.14",
     "@types/pg": "^8.11.8",
+    "@upstash/vector": "^1.1.5",
     "@zilliz/milvus2-sdk-node": "^2.4.6",
     "ajv": "^8.17.1",
     "assemblyai": "^4.7.0",

packages/llamaindex/src/Retriever.ts

@@ -1,20 +0,0 @@
-import type { NodeWithScore } from "@llamaindex/core/schema";
-import type { ServiceContext } from "./ServiceContext.js";
-import type { MessageContent } from "./index.edge.js";
-
-export type RetrieveParams = {
-  query: MessageContent;
-  preFilters?: unknown;
-};
-
-/**
- * Retrievers retrieve the nodes that most closely match our query in similarity.
- */
-export interface BaseRetriever {
-  retrieve(params: RetrieveParams): Promise<NodeWithScore[]>;
-
-  /**
-   * @deprecated to be deprecated soon
-   */
-  serviceContext?: ServiceContext | undefined;
-}

packages/llamaindex/src/agent/anthropic.ts

@@ -1,9 +1,9 @@
-import { Settings } from "../Settings.js";
 import type {
-  ChatEngineParamsNonStreaming,
-  ChatEngineParamsStreaming,
-  EngineResponse,
-} from "../index.edge.js";
+  NonStreamingChatEngineParams,
+  StreamingChatEngineParams,
+} from "@llamaindex/core/chat-engine";
+import type { EngineResponse } from "@llamaindex/core/schema";
+import { Settings } from "../Settings.js";
 import { Anthropic } from "../llm/anthropic.js";
 import { LLMAgent, LLMAgentWorker, type LLMAgentParams } from "./llm.js";
 
@@ -24,12 +24,13 @@ export class AnthropicAgent extends LLMAgent {
     });
   }
 
-  async chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
-  async chat(params: ChatEngineParamsStreaming): Promise<never>;
+  async chat(params: NonStreamingChatEngineParams): Promise<EngineResponse>;
+  async chat(params: StreamingChatEngineParams): Promise<never>;
   override async chat(
-    params: ChatEngineParamsNonStreaming | ChatEngineParamsStreaming,
+    params: NonStreamingChatEngineParams | StreamingChatEngineParams,
   ) {
-    if (params.stream) {
+    const { stream } = params;
+    if (stream) {
       // Anthropic does support this, but looks like it's not supported in the LITS LLM
       throw new Error("Anthropic does not support streaming");
     }

packages/llamaindex/src/agent/base.ts

@@ -1,3 +1,8 @@
+import {
+  BaseChatEngine,
+  type NonStreamingChatEngineParams,
+  type StreamingChatEngineParams,
+} from "@llamaindex/core/chat-engine";
 import type {
   BaseToolWithCall,
   ChatMessage,
@@ -10,11 +15,6 @@ import { EngineResponse } from "@llamaindex/core/schema";
 import { wrapEventCaller } from "@llamaindex/core/utils";
 import { randomUUID } from "@llamaindex/env";
 import { Settings } from "../Settings.js";
-import {
-  type ChatEngine,
-  type ChatEngineParamsNonStreaming,
-  type ChatEngineParamsStreaming,
-} from "../engines/chat/index.js";
 import { consoleLogger, emptyLogger } from "../internal/logger.js";
 import { isReadableStream } from "../internal/utils.js";
 import { ObjectRetriever } from "../objects/index.js";
@@ -207,8 +207,7 @@ export abstract class AgentRunner<
   >
     ? AdditionalMessageOptions
     : never,
-> implements ChatEngine
-{
+> extends BaseChatEngine {
   readonly #llm: AI;
   readonly #tools:
     | BaseToolWithCall[]
@@ -259,6 +258,7 @@ export abstract class AgentRunner<
   protected constructor(
     params: AgentRunnerParams<AI, Store, AdditionalMessageOptions>,
   ) {
+    super();
     const { llm, chatHistory, systemPrompt, runner, tools, verbose } = params;
     this.#llm = llm;
     this.#chatHistory = chatHistory;
@@ -345,20 +345,19 @@ export abstract class AgentRunner<
     });
   }
 
-  async chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
+  async chat(params: NonStreamingChatEngineParams): Promise<EngineResponse>;
   async chat(
-    params: ChatEngineParamsStreaming,
+    params: StreamingChatEngineParams,
   ): Promise<ReadableStream<EngineResponse>>;
   @wrapEventCaller
   async chat(
-    params: ChatEngineParamsNonStreaming | ChatEngineParamsStreaming,
+    params: NonStreamingChatEngineParams | StreamingChatEngineParams,
   ): Promise<EngineResponse | ReadableStream<EngineResponse>> {
     let chatHistory: ChatMessage<AdditionalMessageOptions>[] = [];
 
     if (params.chatHistory instanceof BaseMemory) {
-      chatHistory = (await params.chatHistory.getMessages(
-        params.message,
-      )) as ChatMessage<AdditionalMessageOptions>[];
+      chatHistory =
+        (await params.chatHistory.getMessages()) as ChatMessage<AdditionalMessageOptions>[];
     } else {
       chatHistory =
         params.chatHistory as ChatMessage<AdditionalMessageOptions>[];

packages/llamaindex/src/cloud/LlamaCloudIndex.ts

@@ -1,7 +1,6 @@
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
 import type { Document, TransformComponent } from "@llamaindex/core/schema";
-import type { BaseRetriever } from "../Retriever.js";
 import { RetrieverQueryEngine } from "../engines/query/RetrieverQueryEngine.js";
 import type { BaseNodePostprocessor } from "../postprocessors/types.js";
 import type { CloudRetrieveParams } from "./LlamaCloudRetriever.js";
@@ -12,6 +11,7 @@ import { getAppBaseUrl, getProjectId, initService } from "./utils.js";
 
 import { PipelinesService, ProjectsService } from "@llamaindex/cloud/api";
 import { SentenceSplitter } from "@llamaindex/core/node-parser";
+import type { BaseRetriever } from "@llamaindex/core/retriever";
 import { getEnv } from "@llamaindex/env";
 import { OpenAIEmbedding } from "@llamaindex/openai";
 import { Settings } from "../Settings.js";

packages/llamaindex/src/cloud/LlamaCloudRetriever.ts

@@ -4,11 +4,12 @@ import {
   type RetrievalParams,
   type TextNodeWithScore,
 } from "@llamaindex/cloud/api";
-import { DEFAULT_PROJECT_NAME, Settings } from "@llamaindex/core/global";
+import { DEFAULT_PROJECT_NAME } from "@llamaindex/core/global";
+import type { QueryBundle } from "@llamaindex/core/query-engine";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import type { NodeWithScore } from "@llamaindex/core/schema";
 import { jsonToNode, ObjectType } from "@llamaindex/core/schema";
-import { extractText, wrapEventCaller } from "@llamaindex/core/utils";
-import type { BaseRetriever, RetrieveParams } from "../Retriever.js";
+import { extractText } from "@llamaindex/core/utils";
 import type { ClientParams, CloudConstructorParams } from "./type.js";
 import { getProjectId, initService } from "./utils.js";
 
@@ -17,7 +18,7 @@ export type CloudRetrieveParams = Omit<
   "query" | "search_filters" | "dense_similarity_top_k"
 > & { similarityTopK?: number; filters?: MetadataFilters };
 
-export class LlamaCloudRetriever implements BaseRetriever {
+export class LlamaCloudRetriever extends BaseRetriever {
   clientParams: ClientParams;
   retrieveParams: CloudRetrieveParams;
   organizationId?: string;
@@ -36,12 +37,13 @@ export class LlamaCloudRetriever implements BaseRetriever {
       return {
         // Currently LlamaCloud only supports text nodes
         node: textNode,
-        score: node.score,
+        score: node.score ?? undefined,
       };
     });
   }
 
   constructor(params: CloudConstructorParams & CloudRetrieveParams) {
+    super();
     this.clientParams = { apiKey: params.apiKey, baseUrl: params.baseUrl };
     initService(this.clientParams);
     this.retrieveParams = params;
@@ -54,11 +56,7 @@ export class LlamaCloudRetriever implements BaseRetriever {
     }
   }
 
-  @wrapEventCaller
-  async retrieve({
-    query,
-    preFilters,
-  }: RetrieveParams): Promise<NodeWithScore[]> {
+  async _retrieve(query: QueryBundle): Promise<NodeWithScore[]> {
     const { data: pipelines } =
       await PipelinesService.searchPipelinesApiV1PipelinesGet({
         query: {
@@ -97,19 +95,11 @@ export class LlamaCloudRetriever implements BaseRetriever {
         body: {
           ...this.retrieveParams,
           query: extractText(query),
-          search_filters:
-            this.retrieveParams.filters ?? (preFilters as MetadataFilters),
+          search_filters: this.retrieveParams.filters as MetadataFilters,
           dense_similarity_top_k: this.retrieveParams.similarityTopK!,
         },
       });
 
-    const nodesWithScores = this.resultNodesToNodeWithScore(
-      results.retrieval_nodes,
-    );
-    Settings.callbackManager.dispatchEvent("retrieve-end", {
-      query,
-      nodes: nodesWithScores,
-    });
-    return nodesWithScores;
+    return this.resultNodesToNodeWithScore(results.retrieval_nodes);
   }
 }

packages/llamaindex/src/embeddings/OllamaEmbedding.ts

@@ -1,5 +1,5 @@
 import type { BaseEmbedding } from "@llamaindex/core/embeddings";
-import { Ollama } from "../llm/ollama.js";
+import { Ollama } from "@llamaindex/ollama";
 
 /**
  * OllamaEmbedding is an alias for Ollama that implements the BaseEmbedding interface.

packages/llamaindex/src/engines/chat/CondenseQuestionChatEngine.ts

@@ -1,10 +1,14 @@
+import {
+  BaseChatEngine,
+  type NonStreamingChatEngineParams,
+  type StreamingChatEngineParams,
+} from "@llamaindex/core/chat-engine";
 import type { ChatMessage, LLM } from "@llamaindex/core/llms";
 import { BaseMemory, ChatMemoryBuffer } from "@llamaindex/core/memory";
 import {
   type CondenseQuestionPrompt,
   defaultCondenseQuestionPrompt,
   type ModuleRecord,
-  PromptMixin,
 } from "@llamaindex/core/prompts";
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
 import type { EngineResponse } from "@llamaindex/core/schema";
@@ -16,11 +20,6 @@ import {
 } from "@llamaindex/core/utils";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { llmFromSettingsOrContext } from "../../Settings.js";
-import type {
-  ChatEngine,
-  ChatEngineParamsNonStreaming,
-  ChatEngineParamsStreaming,
-} from "./types.js";
 
 /**
  * CondenseQuestionChatEngine is used in conjunction with a Index (for example VectorStoreIndex).
@@ -32,16 +31,16 @@ import type {
  * underlying data. It performs less well when the chat messages are not questions about the
  * data, or are very referential to previous context.
  */
-
-export class CondenseQuestionChatEngine
-  extends PromptMixin
-  implements ChatEngine
-{
+export class CondenseQuestionChatEngine extends BaseChatEngine {
   queryEngine: BaseQueryEngine;
-  chatHistory: BaseMemory;
+  memory: BaseMemory;
   llm: LLM;
   condenseMessagePrompt: CondenseQuestionPrompt;
 
+  get chatHistory() {
+    return this.memory.getMessages();
+  }
+
   constructor(init: {
     queryEngine: BaseQueryEngine;
     chatHistory: ChatMessage[];
@@ -51,7 +50,7 @@ export class CondenseQuestionChatEngine
     super();
 
     this.queryEngine = init.queryEngine;
-    this.chatHistory = new ChatMemoryBuffer({
+    this.memory = new ChatMemoryBuffer({
       chatHistory: init?.chatHistory,
     });
     this.llm = llmFromSettingsOrContext(init?.serviceContext);
@@ -78,9 +77,7 @@ export class CondenseQuestionChatEngine
   }
 
   private async condenseQuestion(chatHistory: BaseMemory, question: string) {
-    const chatHistoryStr = messagesToHistory(
-      await chatHistory.getMessages(question),
-    );
+    const chatHistoryStr = messagesToHistory(await chatHistory.getMessages());
 
     return this.llm.complete({
       prompt: this.condenseMessagePrompt.format({
@@ -90,23 +87,23 @@ export class CondenseQuestionChatEngine
     });
   }
 
+  chat(params: NonStreamingChatEngineParams): Promise<EngineResponse>;
   chat(
-    params: ChatEngineParamsStreaming,
+    params: StreamingChatEngineParams,
   ): Promise<AsyncIterable<EngineResponse>>;
-  chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
   @wrapEventCaller
   async chat(
-    params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
+    params: NonStreamingChatEngineParams | StreamingChatEngineParams,
   ): Promise<EngineResponse | AsyncIterable<EngineResponse>> {
     const { message, stream } = params;
     const chatHistory = params.chatHistory
       ? new ChatMemoryBuffer({
           chatHistory:
             params.chatHistory instanceof BaseMemory
-              ? await params.chatHistory.getMessages(message)
+              ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
-      : this.chatHistory;
+      : this.memory;
 
     const condensedQuestion = (
       await this.condenseQuestion(chatHistory, extractText(message))
@@ -114,12 +111,10 @@ export class CondenseQuestionChatEngine
     chatHistory.put({ content: message, role: "user" });
 
     if (stream) {
-      const stream = await this.queryEngine.query(
-        {
-          query: condensedQuestion,
-        },
-        true,
-      );
+      const stream = await this.queryEngine.query({
+        query: condensedQuestion,
+        stream: true,
+      });
       return streamReducer({
         stream,
         initialValue: "",
@@ -142,6 +137,6 @@ export class CondenseQuestionChatEngine
   }
 
   reset() {
-    this.chatHistory.reset();
+    this.memory.reset();
   }
 }

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

@@ -1,3 +1,8 @@
+import type {
+  BaseChatEngine,
+  NonStreamingChatEngineParams,
+  StreamingChatEngineParams,
+} from "@llamaindex/core/chat-engine";
 import type {
   ChatMessage,
   LLM,
@@ -11,6 +16,7 @@ import {
   PromptMixin,
   type PromptsRecord,
 } from "@llamaindex/core/prompts";
+import type { BaseRetriever } from "@llamaindex/core/retriever";
 import { EngineResponse, MetadataMode } from "@llamaindex/core/schema";
 import {
   extractText,
@@ -18,27 +24,25 @@ import {
   streamReducer,
   wrapEventCaller,
 } from "@llamaindex/core/utils";
-import type { BaseRetriever } from "../../Retriever.js";
 import { Settings } from "../../Settings.js";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
 import { DefaultContextGenerator } from "./DefaultContextGenerator.js";
-import type {
-  ChatEngine,
-  ChatEngineParamsNonStreaming,
-  ChatEngineParamsStreaming,
-  ContextGenerator,
-} from "./types.js";
+import type { ContextGenerator } from "./types.js";
 
 /**
  * ContextChatEngine uses the Index to get the appropriate context for each query.
  * The context is stored in the system prompt, and the chat history is chunk: ChatResponseChunk, nodes?: NodeWithScore<import("/Users/marcus/code/llamaindex/LlamaIndexTS/packages/core/src/Node").Metadata>[], nodes?: NodeWithScore<import("/Users/marcus/code/llamaindex/LlamaIndexTS/packages/core/src/Node").Metadata>[]lowing the appropriate context to be surfaced for each query.
  */
-export class ContextChatEngine extends PromptMixin implements ChatEngine {
+export class ContextChatEngine extends PromptMixin implements BaseChatEngine {
   chatModel: LLM;
-  chatHistory: BaseMemory;
+  memory: BaseMemory;
   contextGenerator: ContextGenerator & PromptMixin;
   systemPrompt?: string | undefined;
 
+  get chatHistory() {
+    return this.memory.getMessages();
+  }
+
   constructor(init: {
     retriever: BaseRetriever;
     chatModel?: LLM | undefined;
@@ -50,7 +54,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
   }) {
     super();
     this.chatModel = init.chatModel ?? Settings.llm;
-    this.chatHistory = new ChatMemoryBuffer({ chatHistory: init?.chatHistory });
+    this.memory = new ChatMemoryBuffer({ chatHistory: init?.chatHistory });
     this.contextGenerator = new DefaultContextGenerator({
       retriever: init.retriever,
       contextSystemPrompt: init?.contextSystemPrompt,
@@ -79,23 +83,23 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
     };
   }
 
+  chat(params: NonStreamingChatEngineParams): Promise<EngineResponse>;
   chat(
-    params: ChatEngineParamsStreaming,
+    params: StreamingChatEngineParams,
   ): Promise<AsyncIterable<EngineResponse>>;
-  chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
   @wrapEventCaller
   async chat(
-    params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
+    params: StreamingChatEngineParams | NonStreamingChatEngineParams,
   ): Promise<EngineResponse | AsyncIterable<EngineResponse>> {
     const { message, stream } = params;
     const chatHistory = params.chatHistory
       ? new ChatMemoryBuffer({
           chatHistory:
             params.chatHistory instanceof BaseMemory
-              ? await params.chatHistory.getMessages(message)
+              ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
-      : this.chatHistory;
+      : this.memory;
     const requestMessages = await this.prepareRequestMessages(
       message,
       chatHistory,
@@ -125,7 +129,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
   }
 
   reset() {
-    this.chatHistory.reset();
+    this.memory.reset();
   }
 
   private async prepareRequestMessages(
@@ -139,7 +143,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
     const textOnly = extractText(message);
     const context = await this.contextGenerator.generate(textOnly);
     const systemMessage = this.prependSystemPrompt(context.message);
-    const messages = await chatHistory.getMessages(systemMessage.content);
+    const messages = await chatHistory.getMessages([systemMessage]);
     return { nodes: context.nodes, messages };
   }
 

packages/llamaindex/src/engines/chat/DefaultContextGenerator.ts

@@ -5,10 +5,10 @@ import {
   type ModuleRecord,
   PromptMixin,
 } from "@llamaindex/core/prompts";
+import type { BaseRetriever } from "@llamaindex/core/retriever";
 import { MetadataMode, type NodeWithScore } from "@llamaindex/core/schema";
 import { createMessageContent } from "@llamaindex/core/utils";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
-import type { BaseRetriever } from "../../Retriever.js";
 import type { Context, ContextGenerator } from "./types.js";
 
 export class DefaultContextGenerator

packages/llamaindex/src/engines/chat/SimpleChatEngine.ts

@@ -1,3 +1,8 @@
+import type {
+  BaseChatEngine,
+  NonStreamingChatEngineParams,
+  StreamingChatEngineParams,
+} from "@llamaindex/core/chat-engine";
 import type { LLM } from "@llamaindex/core/llms";
 import { BaseMemory, ChatMemoryBuffer } from "@llamaindex/core/memory";
 import { EngineResponse } from "@llamaindex/core/schema";
@@ -7,32 +12,31 @@ import {
   wrapEventCaller,
 } from "@llamaindex/core/utils";
 import { Settings } from "../../Settings.js";
-import type {
-  ChatEngine,
-  ChatEngineParamsNonStreaming,
-  ChatEngineParamsStreaming,
-} from "./types.js";
 
 /**
  * SimpleChatEngine is the simplest possible chat engine. Useful for using your own custom prompts.
  */
 
-export class SimpleChatEngine implements ChatEngine {
-  chatHistory: BaseMemory;
+export class SimpleChatEngine implements BaseChatEngine {
+  memory: BaseMemory;
   llm: LLM;
 
+  get chatHistory() {
+    return this.memory.getMessages();
+  }
+
   constructor(init?: Partial<SimpleChatEngine>) {
-    this.chatHistory = init?.chatHistory ?? new ChatMemoryBuffer();
+    this.memory = init?.memory ?? new ChatMemoryBuffer();
     this.llm = init?.llm ?? Settings.llm;
   }
 
+  chat(params: NonStreamingChatEngineParams): Promise<EngineResponse>;
   chat(
-    params: ChatEngineParamsStreaming,
+    params: StreamingChatEngineParams,
   ): Promise<AsyncIterable<EngineResponse>>;
-  chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
   @wrapEventCaller
   async chat(
-    params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
+    params: NonStreamingChatEngineParams | StreamingChatEngineParams,
   ): Promise<EngineResponse | AsyncIterable<EngineResponse>> {
     const { message, stream } = params;
 
@@ -40,15 +44,15 @@ export class SimpleChatEngine implements ChatEngine {
       ? new ChatMemoryBuffer({
           chatHistory:
             params.chatHistory instanceof BaseMemory
-              ? await params.chatHistory.getMessages(message)
+              ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
-      : this.chatHistory;
+      : this.memory;
     chatHistory.put({ content: message, role: "user" });
 
     if (stream) {
       const stream = await this.llm.chat({
-        messages: await chatHistory.getMessages(params.message),
+        messages: await chatHistory.getMessages(),
         stream: true,
       });
       return streamConverter(
@@ -66,13 +70,13 @@ export class SimpleChatEngine implements ChatEngine {
 
     const response = await this.llm.chat({
       stream: false,
-      messages: await chatHistory.getMessages(params.message),
+      messages: await chatHistory.getMessages(),
     });
     chatHistory.put(response.message);
     return EngineResponse.fromChatResponse(response);
   }
 
   reset() {
-    this.chatHistory.reset();
+    this.memory.reset();
   }
 }

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

@@ -1,58 +1,10 @@
-import type { ChatMessage, MessageContent } from "@llamaindex/core/llms";
-import type { BaseMemory } from "@llamaindex/core/memory";
-import { EngineResponse, type NodeWithScore } from "@llamaindex/core/schema";
-
-/**
- * Represents the base parameters for ChatEngine.
- */
-export interface ChatEngineParamsBase {
-  message: MessageContent;
-  /**
-   * Optional chat history if you want to customize the chat history.
-   */
-  chatHistory?: ChatMessage[] | BaseMemory;
-  /**
-   * Optional flag to enable verbose mode.
-   * @default false
-   */
-  verbose?: boolean;
-}
-
-export interface ChatEngineParamsStreaming extends ChatEngineParamsBase {
-  stream: true;
-}
-
-export interface ChatEngineParamsNonStreaming extends ChatEngineParamsBase {
-  stream?: false | null;
-}
-
-/**
- * A ChatEngine is used to handle back and forth chats between the application and the LLM.
- */
-export interface ChatEngine<
-  // synchronous response
-  R = EngineResponse,
-  // asynchronous response
-  AR extends AsyncIterable<unknown> = AsyncIterable<R>,
-> {
-  /**
-   * Send message along with the class's current chat history to the LLM.
-   * @param params
-   */
-  chat(params: ChatEngineParamsStreaming): Promise<AR>;
-  chat(params: ChatEngineParamsNonStreaming): Promise<R>;
-
-  /**
-   * Resets the chat history so that it's empty.
-   */
-  reset(): void;
-}
+import type { ChatMessage } from "@llamaindex/core/llms";
+import type { NodeWithScore } from "@llamaindex/core/schema";
 
 export interface Context {
   message: ChatMessage;
   nodes: NodeWithScore[];
 }
-
 /**
  * A ContextGenerator is used to generate a context based on a message's text content
  */

packages/llamaindex/src/engines/query/RetrieverQueryEngine.ts

@@ -1,10 +1,11 @@
-import { BaseQueryEngine } from "@llamaindex/core/query-engine";
+import type { MessageContent } from "@llamaindex/core/llms";
+import { BaseQueryEngine, type QueryType } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
 import { getResponseSynthesizer } from "@llamaindex/core/response-synthesizers";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import { type NodeWithScore } from "@llamaindex/core/schema";
 import { extractText } from "@llamaindex/core/utils";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
-import type { BaseRetriever } from "../../Retriever.js";
 
 /**
  * A query engine that uses a retriever to query an index and then synthesizes the response.
@@ -67,7 +68,10 @@ export class RetrieverQueryEngine extends BaseQueryEngine {
     };
   }
 
-  private async applyNodePostprocessors(nodes: NodeWithScore[], query: string) {
+  private async applyNodePostprocessors(
+    nodes: NodeWithScore[],
+    query: MessageContent,
+  ) {
     let nodesWithScore = nodes;
 
     for (const postprocessor of this.nodePostprocessors) {
@@ -80,12 +84,10 @@ export class RetrieverQueryEngine extends BaseQueryEngine {
     return nodesWithScore;
   }
 
-  private async retrieve(query: string) {
-    const nodes = await this.retriever.retrieve({
-      query,
-      preFilters: this.preFilters,
-    });
+  override async retrieve(query: QueryType) {
+    const nodes = await this.retriever.retrieve(query);
 
-    return await this.applyNodePostprocessors(nodes, query);
+    const messageContent = typeof query === "string" ? query : query.query;
+    return await this.applyNodePostprocessors(nodes, messageContent);
   }
 }

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

@@ -136,7 +136,9 @@ export class RouterQueryEngine extends BaseQueryEngine {
         }
 
         const selectedQueryEngine = this.queryEngines[engineInd.index]!;
-        responses.push(await selectedQueryEngine.query(query));
+        responses.push(
+          await selectedQueryEngine.query({ query, stream: false }),
+        );
       }
 
       if (responses.length > 1) {

packages/llamaindex/src/evaluation/Faithfulness.ts

@@ -103,7 +103,8 @@ export class FaithfulnessEvaluator
     });
 
     const responseObj = await queryEngine.query({
-      query: response,
+      query: { query: response },
+      stream: false,
     });
 
     const rawResponseTxt = responseObj.toString();

packages/llamaindex/src/index.edge.ts

@@ -1,6 +1,6 @@
 import type { AgentEndEvent, AgentStartEvent } from "./agent/types.js";
-import type { RetrievalEndEvent, RetrievalStartEvent } from "./llm/types.js";
 
+export * from "@llamaindex/core/chat-engine";
 export {
   CallbackManager,
   DEFAULT_BASE_URL,
@@ -35,12 +35,11 @@ export * from "@llamaindex/core/llms";
 export * from "@llamaindex/core/prompts";
 export * from "@llamaindex/core/query-engine";
 export * from "@llamaindex/core/response-synthesizers";
+export * from "@llamaindex/core/retriever";
 export * from "@llamaindex/core/schema";
 
 declare module "@llamaindex/core/global" {
   export interface LlamaIndexEventMaps {
-    "retrieve-start": RetrievalStartEvent;
-    "retrieve-end": RetrievalEndEvent;
     // agent events
     "agent-start": AgentStartEvent;
     "agent-end": AgentEndEvent;
@@ -66,7 +65,6 @@ export * from "./objects/index.js";
 export * from "./OutputParser.js";
 export * from "./postprocessors/index.js";
 export * from "./QuestionGenerator.js";
-export * from "./Retriever.js";
 export * from "./selectors/index.js";
 export * from "./ServiceContext.js";
 export { Settings } from "./Settings.js";

packages/llamaindex/src/indices/BaseIndex.ts

@@ -1,7 +1,7 @@
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
+import type { BaseRetriever } from "@llamaindex/core/retriever";
 import type { BaseNode, Document } from "@llamaindex/core/schema";
-import type { BaseRetriever } from "../Retriever.js";
 import type { ServiceContext } from "../ServiceContext.js";
 import { nodeParserFromSettingsOrContext } from "../Settings.js";
 import { runTransformations } from "../ingestion/IngestionPipeline.js";

packages/llamaindex/src/indices/keyword/index.ts

@@ -5,7 +5,6 @@ import type {
   NodeWithScore,
 } from "@llamaindex/core/schema";
 import { MetadataMode } from "@llamaindex/core/schema";
-import type { BaseRetriever, RetrieveParams } from "../../Retriever.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { serviceContextFromDefaults } from "../../ServiceContext.js";
 import { RetrieverQueryEngine } from "../../engines/query/index.js";
@@ -29,7 +28,11 @@ import {
   type KeywordExtractPrompt,
   type QueryKeywordExtractPrompt,
 } from "@llamaindex/core/prompts";
-import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
+import type {
+  BaseQueryEngine,
+  QueryBundle,
+} from "@llamaindex/core/query-engine";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import { extractText } from "@llamaindex/core/utils";
 import { llmFromSettingsOrContext } from "../../Settings.js";
 
@@ -48,7 +51,7 @@ export enum KeywordTableRetrieverMode {
 }
 
 // Base Keyword Table Retriever
-abstract class BaseKeywordTableRetriever implements BaseRetriever {
+abstract class BaseKeywordTableRetriever extends BaseRetriever {
   protected index: KeywordTableIndex;
   protected indexStruct: KeywordTable;
   protected docstore: BaseDocumentStore;
@@ -72,6 +75,7 @@ abstract class BaseKeywordTableRetriever implements BaseRetriever {
     maxKeywordsPerQuery: number;
     numChunksPerQuery: number;
   }) {
+    super();
     this.index = index;
     this.indexStruct = index.indexStruct;
     this.docstore = index.docStore;
@@ -87,7 +91,7 @@ abstract class BaseKeywordTableRetriever implements BaseRetriever {
 
   abstract getKeywords(query: string): Promise<string[]>;
 
-  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
+  async _retrieve(query: QueryBundle): Promise<NodeWithScore[]> {
     const keywords = await this.getKeywords(extractText(query));
     const chunkIndicesCount: { [key: string]: number } = {};
     const filteredKeywords = keywords.filter((keyword) =>

packages/llamaindex/src/indices/summary/index.ts

@@ -2,16 +2,17 @@ import {
   type ChoiceSelectPrompt,
   defaultChoiceSelectPrompt,
 } from "@llamaindex/core/prompts";
+import type { QueryBundle } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
 import { getResponseSynthesizer } from "@llamaindex/core/response-synthesizers";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import type {
   BaseNode,
   Document,
   NodeWithScore,
 } from "@llamaindex/core/schema";
-import { extractText, wrapEventCaller } from "@llamaindex/core/utils";
+import { extractText } from "@llamaindex/core/utils";
 import _ from "lodash";
-import type { BaseRetriever, RetrieveParams } from "../../Retriever.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import {
   llmFromSettingsOrContext,
@@ -279,15 +280,15 @@ export type ListRetrieverMode = SummaryRetrieverMode;
 /**
  * Simple retriever for SummaryIndex that returns all nodes
  */
-export class SummaryIndexRetriever implements BaseRetriever {
+export class SummaryIndexRetriever extends BaseRetriever {
   index: SummaryIndex;
 
   constructor(index: SummaryIndex) {
+    super();
     this.index = index;
   }
 
-  @wrapEventCaller
-  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
+  async _retrieve(queryBundle: QueryBundle): Promise<NodeWithScore[]> {
     const nodeIds = this.index.indexStruct.nodes;
     const nodes = await this.index.docStore.getNodes(nodeIds);
     return nodes.map((node) => ({
@@ -300,7 +301,7 @@ export class SummaryIndexRetriever implements BaseRetriever {
 /**
  * LLM retriever for SummaryIndex which lets you select the most relevant chunks.
  */
-export class SummaryIndexLLMRetriever implements BaseRetriever {
+export class SummaryIndexLLMRetriever extends BaseRetriever {
   index: SummaryIndex;
   choiceSelectPrompt: ChoiceSelectPrompt;
   choiceBatchSize: number;
@@ -317,6 +318,7 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
     parseChoiceSelectAnswerFn?: ChoiceSelectParserFunction,
     serviceContext?: ServiceContext,
   ) {
+    super();
     this.index = index;
     this.choiceSelectPrompt = choiceSelectPrompt || defaultChoiceSelectPrompt;
     this.choiceBatchSize = choiceBatchSize;
@@ -326,7 +328,7 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
     this.serviceContext = serviceContext || index.serviceContext;
   }
 
-  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
+  async _retrieve(query: QueryBundle): Promise<NodeWithScore[]> {
     const nodeIds = this.index.indexStruct.nodes;
     const results: NodeWithScore[] = [];
 

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

@@ -2,9 +2,10 @@ import {
   DEFAULT_SIMILARITY_TOP_K,
   type BaseEmbedding,
 } from "@llamaindex/core/embeddings";
-import { Settings } from "@llamaindex/core/global";
 import type { MessageContent } from "@llamaindex/core/llms";
+import type { QueryBundle } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import {
   ImageNode,
   ModalityType,
@@ -14,8 +15,6 @@ import {
   type Document,
   type NodeWithScore,
 } from "@llamaindex/core/schema";
-import { wrapEventCaller } from "@llamaindex/core/utils";
-import type { BaseRetriever, RetrieveParams } from "../../Retriever.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { nodeParserFromSettingsOrContext } from "../../Settings.js";
 import { RetrieverQueryEngine } from "../../engines/query/RetrieverQueryEngine.js";
@@ -388,7 +387,7 @@ export type VectorIndexRetrieverOptions = {
   filters?: MetadataFilters;
 };
 
-export class VectorIndexRetriever implements BaseRetriever {
+export class VectorIndexRetriever extends BaseRetriever {
   index: VectorStoreIndex;
   topK: TopKMap;
 
@@ -401,6 +400,7 @@ export class VectorIndexRetriever implements BaseRetriever {
     topK,
     filters,
   }: VectorIndexRetrieverOptions) {
+    super();
     this.index = index;
     this.serviceContext = this.index.serviceContext;
     this.topK = topK ?? {
@@ -417,32 +417,17 @@ export class VectorIndexRetriever implements BaseRetriever {
     this.topK[ModalityType.TEXT] = similarityTopK;
   }
 
-  @wrapEventCaller
-  async retrieve({
-    query,
-    preFilters,
-  }: RetrieveParams): Promise<NodeWithScore[]> {
-    Settings.callbackManager.dispatchEvent("retrieve-start", {
-      query,
-    });
+  async _retrieve(params: QueryBundle): Promise<NodeWithScore[]> {
+    const { query } = params;
     const vectorStores = this.index.vectorStores;
     let nodesWithScores: NodeWithScore[] = [];
 
     for (const type in vectorStores) {
       const vectorStore: VectorStore = vectorStores[type as ModalityType]!;
       nodesWithScores = nodesWithScores.concat(
-        await this.retrieveQuery(
-          query,
-          type as ModalityType,
-          vectorStore,
-          preFilters as MetadataFilters,
-        ),
+        await this.retrieveQuery(query, type as ModalityType, vectorStore),
       );
     }
-    Settings.callbackManager.dispatchEvent("retrieve-end", {
-      query,
-      nodes: nodesWithScores,
-    });
     return nodesWithScores;
   }
 

packages/llamaindex/src/ingestion/IngestionCache.ts

@@ -1,7 +1,11 @@
 import type { BaseNode, TransformComponent } from "@llamaindex/core/schema";
 import { MetadataMode } from "@llamaindex/core/schema";
 import { createSHA256 } from "@llamaindex/env";
-import { docToJson, jsonToDoc } from "../storage/docStore/utils.js";
+import {
+  docToJson,
+  jsonSerializer,
+  jsonToDoc,
+} from "../storage/docStore/utils.js";
 import { SimpleKVStore } from "../storage/kvStore/SimpleKVStore.js";
 import type { BaseKVStore } from "../storage/kvStore/types.js";
 
@@ -53,7 +57,7 @@ export class IngestionCache {
 
   async put(hash: string, nodes: BaseNode[]) {
     const val = {
-      [this.nodesKey]: nodes.map((node) => docToJson(node)),
+      [this.nodesKey]: nodes.map((node) => docToJson(node, jsonSerializer)),
     };
     await this.cache.put(hash, val, this.collection);
   }
@@ -63,6 +67,8 @@ export class IngestionCache {
     if (!json || !json[this.nodesKey] || !Array.isArray(json[this.nodesKey])) {
       return undefined;
     }
-    return json[this.nodesKey].map((doc: any) => jsonToDoc(doc));
+    return json[this.nodesKey].map((doc: any) =>
+      jsonToDoc(doc, jsonSerializer),
+    );
   }
 }

packages/llamaindex/src/internal/deps/ollama.d.ts

@@ -1,264 +0,0 @@
-type Fetch = typeof fetch;
-interface Config {
-  host: string;
-  fetch?: Fetch;
-  proxy?: boolean;
-}
-interface Options {
-  numa: boolean;
-  num_ctx: number;
-  num_batch: number;
-  main_gpu: number;
-  low_vram: boolean;
-  f16_kv: boolean;
-  logits_all: boolean;
-  vocab_only: boolean;
-  use_mmap: boolean;
-  use_mlock: boolean;
-  embedding_only: boolean;
-  num_thread: number;
-  num_keep: number;
-  seed: number;
-  num_predict: number;
-  top_k: number;
-  top_p: number;
-  tfs_z: number;
-  typical_p: number;
-  repeat_last_n: number;
-  temperature: number;
-  repeat_penalty: number;
-  presence_penalty: number;
-  frequency_penalty: number;
-  mirostat: number;
-  mirostat_tau: number;
-  mirostat_eta: number;
-  penalize_newline: boolean;
-  stop: string[];
-}
-interface GenerateRequest {
-  model: string;
-  prompt: string;
-  system?: string;
-  template?: string;
-  context?: number[];
-  stream?: boolean;
-  raw?: boolean;
-  format?: string;
-  images?: Uint8Array[] | string[];
-  keep_alive?: string | number;
-  options?: Partial<Options>;
-}
-interface Message {
-  role: string;
-  content: string;
-  images?: Uint8Array[] | string[];
-}
-interface ChatRequest {
-  model: string;
-  messages?: Message[];
-  stream?: boolean;
-  format?: string;
-  keep_alive?: string | number;
-  options?: Partial<Options>;
-}
-interface PullRequest {
-  model: string;
-  insecure?: boolean;
-  stream?: boolean;
-}
-interface PushRequest {
-  model: string;
-  insecure?: boolean;
-  stream?: boolean;
-}
-interface CreateRequest {
-  model: string;
-  path?: string;
-  modelfile?: string;
-  stream?: boolean;
-}
-interface DeleteRequest {
-  model: string;
-}
-interface CopyRequest {
-  source: string;
-  destination: string;
-}
-interface ShowRequest {
-  model: string;
-  system?: string;
-  template?: string;
-  options?: Partial<Options>;
-}
-interface EmbeddingsRequest {
-  model: string;
-  prompt: string;
-  keep_alive?: string | number;
-  options?: Partial<Options>;
-}
-interface GenerateResponse {
-  model: string;
-  created_at: Date;
-  response: string;
-  done: boolean;
-  context: number[];
-  total_duration: number;
-  load_duration: number;
-  prompt_eval_count: number;
-  prompt_eval_duration: number;
-  eval_count: number;
-  eval_duration: number;
-}
-interface ChatResponse {
-  model: string;
-  created_at: Date;
-  message: Message;
-  done: boolean;
-  total_duration: number;
-  load_duration: number;
-  prompt_eval_count: number;
-  prompt_eval_duration: number;
-  eval_count: number;
-  eval_duration: number;
-}
-interface EmbeddingsResponse {
-  embedding: number[];
-}
-interface ProgressResponse {
-  status: string;
-  digest: string;
-  total: number;
-  completed: number;
-}
-interface ModelResponse {
-  name: string;
-  modified_at: Date;
-  size: number;
-  digest: string;
-  details: ModelDetails;
-}
-interface ModelDetails {
-  parent_model: string;
-  format: string;
-  family: string;
-  families: string[];
-  parameter_size: string;
-  quantization_level: string;
-}
-interface ShowResponse {
-  license: string;
-  modelfile: string;
-  parameters: string;
-  template: string;
-  system: string;
-  details: ModelDetails;
-  messages: Message[];
-}
-interface ListResponse {
-  models: ModelResponse[];
-}
-interface ErrorResponse {
-  error: string;
-}
-interface StatusResponse {
-  status: string;
-}
-
-declare class Ollama {
-  protected readonly config: Config;
-  protected readonly fetch: Fetch;
-  private abortController;
-  constructor(config?: Partial<Config>);
-  abort(): void;
-  protected processStreamableRequest<T extends object>(
-    endpoint: string,
-    request: {
-      stream?: boolean;
-    } & Record<string, any>,
-  ): Promise<T | AsyncGenerator<T>>;
-  encodeImage(image: Uint8Array | string): Promise<string>;
-  generate(
-    request: GenerateRequest & {
-      stream: true;
-    },
-  ): Promise<AsyncGenerator<GenerateResponse>>;
-  generate(
-    request: GenerateRequest & {
-      stream?: false;
-    },
-  ): Promise<GenerateResponse>;
-  chat(
-    request: ChatRequest & {
-      stream: true;
-    },
-  ): Promise<AsyncGenerator<ChatResponse>>;
-  chat(
-    request: ChatRequest & {
-      stream?: false;
-    },
-  ): Promise<ChatResponse>;
-  create(
-    request: CreateRequest & {
-      stream: true;
-    },
-  ): Promise<AsyncGenerator<ProgressResponse>>;
-  create(
-    request: CreateRequest & {
-      stream?: false;
-    },
-  ): Promise<ProgressResponse>;
-  pull(
-    request: PullRequest & {
-      stream: true;
-    },
-  ): Promise<AsyncGenerator<ProgressResponse>>;
-  pull(
-    request: PullRequest & {
-      stream?: false;
-    },
-  ): Promise<ProgressResponse>;
-  push(
-    request: PushRequest & {
-      stream: true;
-    },
-  ): Promise<AsyncGenerator<ProgressResponse>>;
-  push(
-    request: PushRequest & {
-      stream?: false;
-    },
-  ): Promise<ProgressResponse>;
-  delete(request: DeleteRequest): Promise<StatusResponse>;
-  copy(request: CopyRequest): Promise<StatusResponse>;
-  list(): Promise<ListResponse>;
-  show(request: ShowRequest): Promise<ShowResponse>;
-  embeddings(request: EmbeddingsRequest): Promise<EmbeddingsResponse>;
-}
-declare const _default: Ollama;
-
-export {
-  Ollama,
-  _default as default,
-  type ChatRequest,
-  type ChatResponse,
-  type Config,
-  type CopyRequest,
-  type CreateRequest,
-  type DeleteRequest,
-  type EmbeddingsRequest,
-  type EmbeddingsResponse,
-  type ErrorResponse,
-  type Fetch,
-  type GenerateRequest,
-  type GenerateResponse,
-  type ListResponse,
-  type Message,
-  type ModelDetails,
-  type ModelResponse,
-  type Options,
-  type ProgressResponse,
-  type PullRequest,
-  type PushRequest,
-  type ShowRequest,
-  type ShowResponse,
-  type StatusResponse,
-};

packages/llamaindex/src/internal/deps/ollama.js

@@ -1,462 +0,0 @@
-// generate from "tsup ./src/browser.js --format esm --dts"
-var __defProp = Object.defineProperty;
-var __getOwnPropSymbols = Object.getOwnPropertySymbols;
-var __hasOwnProp = Object.prototype.hasOwnProperty;
-var __propIsEnum = Object.prototype.propertyIsEnumerable;
-var __knownSymbol = (name, symbol) => {
-  return (symbol = Symbol[name]) ? symbol : Symbol.for("Symbol." + name);
-};
-var __defNormalProp = (obj, key, value) =>
-  key in obj
-    ? __defProp(obj, key, {
-        enumerable: true,
-        configurable: true,
-        writable: true,
-        value,
-      })
-    : (obj[key] = value);
-var __spreadValues = (a, b) => {
-  for (var prop in b || (b = {}))
-    if (__hasOwnProp.call(b, prop)) __defNormalProp(a, prop, b[prop]);
-  if (__getOwnPropSymbols)
-    for (var prop of __getOwnPropSymbols(b)) {
-      if (__propIsEnum.call(b, prop)) __defNormalProp(a, prop, b[prop]);
-    }
-  return a;
-};
-var __async = (__this, __arguments, generator) => {
-  return new Promise((resolve, reject) => {
-    var fulfilled = (value) => {
-      try {
-        step(generator.next(value));
-      } catch (e) {
-        reject(e);
-      }
-    };
-    var rejected = (value) => {
-      try {
-        step(generator.throw(value));
-      } catch (e) {
-        reject(e);
-      }
-    };
-    var step = (x) =>
-      x.done
-        ? resolve(x.value)
-        : Promise.resolve(x.value).then(fulfilled, rejected);
-    step((generator = generator.apply(__this, __arguments)).next());
-  });
-};
-var __await = function (promise, isYieldStar) {
-  this[0] = promise;
-  this[1] = isYieldStar;
-};
-var __asyncGenerator = (__this, __arguments, generator) => {
-  var resume = (k, v, yes, no) => {
-    try {
-      var x = generator[k](v),
-        isAwait = (v = x.value) instanceof __await,
-        done = x.done;
-      Promise.resolve(isAwait ? v[0] : v)
-        .then((y) =>
-          isAwait
-            ? resume(
-                k === "return" ? k : "next",
-                v[1] ? { done: y.done, value: y.value } : y,
-                yes,
-                no,
-              )
-            : yes({ value: y, done }),
-        )
-        .catch((e) => resume("throw", e, yes, no));
-    } catch (e) {
-      no(e);
-    }
-  };
-  var method = (k) =>
-    (it[k] = (x) => new Promise((yes, no) => resume(k, x, yes, no)));
-  var it = {};
-  return (
-    (generator = generator.apply(__this, __arguments)),
-    (it[__knownSymbol("asyncIterator")] = () => it),
-    method("next"),
-    method("throw"),
-    method("return"),
-    it
-  );
-};
-var __forAwait = (obj, it, method) =>
-  (it = obj[__knownSymbol("asyncIterator")])
-    ? it.call(obj)
-    : ((obj = obj[__knownSymbol("iterator")]()),
-      (it = {}),
-      (method = (key, fn) =>
-        (fn = obj[key]) &&
-        (it[key] = (arg) =>
-          new Promise(
-            (yes, no, done) => (
-              (arg = fn.call(obj, arg)),
-              (done = arg.done),
-              Promise.resolve(arg.value).then(
-                (value) => yes({ value, done }),
-                no,
-              )
-            ),
-          ))),
-      method("next"),
-      method("return"),
-      it);
-
-// src/version.ts
-var version = "0.0.0";
-
-// src/utils.ts
-var ResponseError = class _ResponseError extends Error {
-  constructor(error, status_code) {
-    super(error);
-    this.error = error;
-    this.status_code = status_code;
-    this.name = "ResponseError";
-    if (Error.captureStackTrace) {
-      Error.captureStackTrace(this, _ResponseError);
-    }
-  }
-};
-var checkOk = (response) =>
-  __async(void 0, null, function* () {
-    var _a;
-    if (!response.ok) {
-      let message = `Error ${response.status}: ${response.statusText}`;
-      let errorData = null;
-      if (
-        (_a = response.headers.get("content-type")) == null
-          ? void 0
-          : _a.includes("application/json")
-      ) {
-        try {
-          errorData = yield response.json();
-          message = errorData.error || message;
-        } catch (error) {
-          console.log("Failed to parse error response as JSON");
-        }
-      } else {
-        try {
-          console.log("Getting text from response");
-          const textResponse = yield response.text();
-          message = textResponse || message;
-        } catch (error) {
-          console.log("Failed to get text from error response");
-        }
-      }
-      throw new ResponseError(message, response.status);
-    }
-  });
-function getPlatform() {
-  if (typeof window !== "undefined" && window.navigator) {
-    return `${window.navigator.platform.toLowerCase()} Browser/${navigator.userAgent};`;
-  } else if (typeof process !== "undefined") {
-    return `${process.arch} ${process.platform} Node.js/${process.version}`;
-  }
-  return "";
-}
-var fetchWithHeaders = (_0, _1, ..._2) =>
-  __async(void 0, [_0, _1, ..._2], function* (fetch2, url, options = {}) {
-    const defaultHeaders = {
-      "Content-Type": "application/json",
-      Accept: "application/json",
-      "User-Agent": `ollama-js/${version} (${getPlatform()})`,
-    };
-    if (!options.headers) {
-      options.headers = {};
-    }
-    options.headers = __spreadValues(
-      __spreadValues({}, defaultHeaders),
-      options.headers,
-    );
-    return fetch2(url, options);
-  });
-var get = (fetch2, host) =>
-  __async(void 0, null, function* () {
-    const response = yield fetchWithHeaders(fetch2, host);
-    yield checkOk(response);
-    return response;
-  });
-var post = (fetch2, host, data, options) =>
-  __async(void 0, null, function* () {
-    const isRecord = (input) => {
-      return (
-        input !== null && typeof input === "object" && !Array.isArray(input)
-      );
-    };
-    const formattedData = isRecord(data) ? JSON.stringify(data) : data;
-    const response = yield fetchWithHeaders(fetch2, host, {
-      method: "POST",
-      body: formattedData,
-      signal: options == null ? void 0 : options.signal,
-    });
-    yield checkOk(response);
-    return response;
-  });
-var del = (fetch2, host, data) =>
-  __async(void 0, null, function* () {
-    const response = yield fetchWithHeaders(fetch2, host, {
-      method: "DELETE",
-      body: JSON.stringify(data),
-    });
-    yield checkOk(response);
-    return response;
-  });
-var parseJSON = function (itr) {
-  return __asyncGenerator(this, null, function* () {
-    var _a;
-    const decoder = new TextDecoder("utf-8");
-    let buffer = "";
-    const reader = itr.getReader();
-    while (true) {
-      const { done, value: chunk } = yield new __await(reader.read());
-      if (done) {
-        break;
-      }
-      buffer += decoder.decode(chunk);
-      const parts = buffer.split("\n");
-      buffer = (_a = parts.pop()) != null ? _a : "";
-      for (const part of parts) {
-        try {
-          yield JSON.parse(part);
-        } catch (error) {
-          console.warn("invalid json: ", part);
-        }
-      }
-    }
-    for (const part of buffer.split("\n").filter((p) => p !== "")) {
-      try {
-        yield JSON.parse(part);
-      } catch (error) {
-        console.warn("invalid json: ", part);
-      }
-    }
-  });
-};
-var formatHost = (host) => {
-  if (!host) {
-    return "http://127.0.0.1:11434";
-  }
-  let isExplicitProtocol = host.includes("://");
-  if (host.startsWith(":")) {
-    host = `http://127.0.0.1${host}`;
-    isExplicitProtocol = false;
-  }
-  if (!isExplicitProtocol) {
-    host = `http://${host}`;
-  }
-  const url = new URL(host);
-  let port = url.port;
-  if (!port) {
-    if (!isExplicitProtocol) {
-      port = "11434";
-    } else {
-      port = url.protocol === "https:" ? "443" : "80";
-    }
-  }
-  let formattedHost = `${url.protocol}//${url.hostname}:${port}${url.pathname}`;
-  if (formattedHost.endsWith("/")) {
-    formattedHost = formattedHost.slice(0, -1);
-  }
-  return formattedHost;
-};
-
-// src/browser.ts
-// import "whatwg-fetch";
-var Ollama = class {
-  constructor(config) {
-    var _a;
-    this.config = {
-      host: "",
-    };
-    if (!(config == null ? void 0 : config.proxy)) {
-      this.config.host = formatHost(
-        (_a = config == null ? void 0 : config.host) != null
-          ? _a
-          : "http://127.0.0.1:11434",
-      );
-    }
-    this.fetch = fetch;
-    if ((config == null ? void 0 : config.fetch) != null) {
-      this.fetch = config.fetch;
-    }
-    this.abortController = new AbortController();
-  }
-  // Abort any ongoing requests to Ollama
-  abort() {
-    this.abortController.abort();
-    this.abortController = new AbortController();
-  }
-  processStreamableRequest(endpoint, request) {
-    return __async(this, null, function* () {
-      var _a;
-      request.stream = (_a = request.stream) != null ? _a : false;
-      const response = yield post(
-        this.fetch,
-        `${this.config.host}/api/${endpoint}`,
-        __spreadValues({}, request),
-        { signal: this.abortController.signal },
-      );
-      if (!response.body) {
-        throw new Error("Missing body");
-      }
-      const itr = parseJSON(response.body);
-      if (request.stream) {
-        return (function () {
-          return __asyncGenerator(this, null, function* () {
-            try {
-              for (
-                var iter = __forAwait(itr), more, temp, error;
-                (more = !(temp = yield new __await(iter.next())).done);
-                more = false
-              ) {
-                const message = temp.value;
-                if ("error" in message) {
-                  throw new Error(message.error);
-                }
-                yield message;
-                if (message.done || message.status === "success") {
-                  return;
-                }
-              }
-            } catch (temp) {
-              error = [temp];
-            } finally {
-              try {
-                more &&
-                  (temp = iter.return) &&
-                  (yield new __await(temp.call(iter)));
-              } finally {
-                if (error) throw error[0];
-              }
-            }
-            throw new Error(
-              "Did not receive done or success response in stream.",
-            );
-          });
-        })();
-      } else {
-        const message = yield itr.next();
-        if (!message.value.done && message.value.status !== "success") {
-          throw new Error("Expected a completed response.");
-        }
-        return message.value;
-      }
-    });
-  }
-  encodeImage(image) {
-    return __async(this, null, function* () {
-      if (typeof image !== "string") {
-        const uint8Array = new Uint8Array(image);
-        const numberArray = Array.from(uint8Array);
-        const base64String = btoa(String.fromCharCode.apply(null, numberArray));
-        return base64String;
-      }
-      return image;
-    });
-  }
-  generate(request) {
-    return __async(this, null, function* () {
-      if (request.images) {
-        request.images = yield Promise.all(
-          request.images.map(this.encodeImage.bind(this)),
-        );
-      }
-      return this.processStreamableRequest("generate", request);
-    });
-  }
-  chat(request) {
-    return __async(this, null, function* () {
-      if (request.messages) {
-        for (const message of request.messages) {
-          if (message.images) {
-            message.images = yield Promise.all(
-              message.images.map(this.encodeImage.bind(this)),
-            );
-          }
-        }
-      }
-      return this.processStreamableRequest("chat", request);
-    });
-  }
-  create(request) {
-    return __async(this, null, function* () {
-      return this.processStreamableRequest("create", {
-        name: request.model,
-        stream: request.stream,
-        modelfile: request.modelfile,
-      });
-    });
-  }
-  pull(request) {
-    return __async(this, null, function* () {
-      return this.processStreamableRequest("pull", {
-        name: request.model,
-        stream: request.stream,
-        insecure: request.insecure,
-      });
-    });
-  }
-  push(request) {
-    return __async(this, null, function* () {
-      return this.processStreamableRequest("push", {
-        name: request.model,
-        stream: request.stream,
-        insecure: request.insecure,
-      });
-    });
-  }
-  delete(request) {
-    return __async(this, null, function* () {
-      yield del(this.fetch, `${this.config.host}/api/delete`, {
-        name: request.model,
-      });
-      return { status: "success" };
-    });
-  }
-  copy(request) {
-    return __async(this, null, function* () {
-      yield post(
-        this.fetch,
-        `${this.config.host}/api/copy`,
-        __spreadValues({}, request),
-      );
-      return { status: "success" };
-    });
-  }
-  list() {
-    return __async(this, null, function* () {
-      const response = yield get(this.fetch, `${this.config.host}/api/tags`);
-      const listResponse = yield response.json();
-      return listResponse;
-    });
-  }
-  show(request) {
-    return __async(this, null, function* () {
-      const response = yield post(
-        this.fetch,
-        `${this.config.host}/api/show`,
-        __spreadValues({}, request),
-      );
-      const showResponse = yield response.json();
-      return showResponse;
-    });
-  }
-  embeddings(request) {
-    return __async(this, null, function* () {
-      const response = yield post(
-        this.fetch,
-        `${this.config.host}/api/embeddings`,
-        __spreadValues({}, request),
-      );
-      const embeddingsResponse = yield response.json();
-      return embeddingsResponse;
-    });
-  }
-};
-var browser_default = new Ollama();
-export { Ollama, browser_default as default };

packages/llamaindex/src/internal/deps/ollama.license

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2023 Saul
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

packages/llamaindex/src/llm/index.ts

@@ -23,7 +23,7 @@ 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 { DeepInfra } from "./deepinfra.js";
-export { Ollama, type OllamaParams } from "./ollama.js";
+export * from "./ollama.js";
 export {
   ALL_AVAILABLE_REPLICATE_MODELS,
   DeuceChatStrategy,

packages/llamaindex/src/llm/ollama.ts

@@ -1,253 +1 @@
-import { BaseEmbedding } from "@llamaindex/core/embeddings";
-import type {
-  ChatResponse,
-  ChatResponseChunk,
-  CompletionResponse,
-  LLM,
-  LLMChatParamsNonStreaming,
-  LLMChatParamsStreaming,
-  LLMCompletionParamsNonStreaming,
-  LLMCompletionParamsStreaming,
-  LLMMetadata,
-} from "@llamaindex/core/llms";
-import { extractText, streamConverter } from "@llamaindex/core/utils";
-import {
-  Ollama as OllamaBase,
-  type Config,
-  type CopyRequest,
-  type CreateRequest,
-  type DeleteRequest,
-  type EmbeddingsRequest,
-  type EmbeddingsResponse,
-  type GenerateRequest,
-  type ListResponse,
-  type ChatResponse as OllamaChatResponse,
-  type GenerateResponse as OllamaGenerateResponse,
-  type Options,
-  type ProgressResponse,
-  type PullRequest,
-  type PushRequest,
-  type ShowRequest,
-  type ShowResponse,
-  type StatusResponse,
-} from "../internal/deps/ollama.js";
-
-const messageAccessor = (part: OllamaChatResponse): ChatResponseChunk => {
-  return {
-    raw: part,
-    delta: part.message.content,
-  };
-};
-
-const completionAccessor = (
-  part: OllamaGenerateResponse,
-): CompletionResponse => {
-  return { text: part.response, raw: part };
-};
-
-export type OllamaParams = {
-  model: string;
-  config?: Partial<Config>;
-  options?: Partial<Options>;
-};
-
-/**
- * This class both implements the LLM and Embedding interfaces.
- */
-export class Ollama
-  extends BaseEmbedding
-  implements LLM, Omit<OllamaBase, "chat">
-{
-  readonly hasStreaming = true;
-
-  ollama: OllamaBase;
-
-  // https://ollama.ai/library
-  model: string;
-
-  options: Partial<Omit<Options, "num_ctx" | "top_p" | "temperature">> &
-    Pick<Options, "num_ctx" | "top_p" | "temperature"> = {
-    num_ctx: 4096,
-    top_p: 0.9,
-    temperature: 0.7,
-  };
-
-  constructor(params: OllamaParams) {
-    super();
-    this.model = params.model;
-    this.ollama = new OllamaBase(params.config);
-    if (params.options) {
-      this.options = {
-        ...this.options,
-        ...params.options,
-      };
-    }
-  }
-
-  get metadata(): LLMMetadata {
-    const { temperature, top_p, num_ctx } = this.options;
-    return {
-      model: this.model,
-      temperature: temperature,
-      topP: top_p,
-      maxTokens: undefined,
-      contextWindow: num_ctx,
-      tokenizer: undefined,
-    };
-  }
-
-  chat(
-    params: LLMChatParamsStreaming,
-  ): Promise<AsyncIterable<ChatResponseChunk>>;
-  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
-  async chat(
-    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
-  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, stream } = params;
-    const payload = {
-      model: this.model,
-      messages: messages.map((message) => ({
-        role: message.role,
-        content: extractText(message.content),
-      })),
-      stream: !!stream,
-      options: {
-        ...this.options,
-      },
-    };
-    if (!stream) {
-      const chatResponse = await this.ollama.chat({
-        ...payload,
-        stream: false,
-      });
-
-      return {
-        message: {
-          role: "assistant",
-          content: chatResponse.message.content,
-        },
-        raw: chatResponse,
-      };
-    } else {
-      const stream = await this.ollama.chat({
-        ...payload,
-        stream: true,
-      });
-      return streamConverter(stream, messageAccessor);
-    }
-  }
-
-  complete(
-    params: LLMCompletionParamsStreaming,
-  ): Promise<AsyncIterable<CompletionResponse>>;
-  complete(
-    params: LLMCompletionParamsNonStreaming,
-  ): Promise<CompletionResponse>;
-  async complete(
-    params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
-  ): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
-    const { prompt, stream } = params;
-    const payload = {
-      model: this.model,
-      prompt: extractText(prompt),
-      stream: !!stream,
-      options: {
-        ...this.options,
-      },
-    };
-    if (!stream) {
-      const response = await this.ollama.generate({
-        ...payload,
-        stream: false,
-      });
-      return {
-        text: response.response,
-        raw: response,
-      };
-    } else {
-      const stream = await this.ollama.generate({
-        ...payload,
-        stream: true,
-      });
-      return streamConverter(stream, completionAccessor);
-    }
-  }
-
-  private async getEmbedding(prompt: string): Promise<number[]> {
-    const payload = {
-      model: this.model,
-      prompt,
-      options: {
-        ...this.options,
-      },
-    };
-    const response = await this.ollama.embeddings({
-      ...payload,
-    });
-    return response.embedding;
-  }
-
-  async getTextEmbedding(text: string): Promise<number[]> {
-    return this.getEmbedding(text);
-  }
-
-  // Inherited from OllamaBase
-
-  push(
-    request: PushRequest & { stream: true },
-  ): Promise<AsyncGenerator<ProgressResponse, any, unknown>>;
-  push(
-    request: PushRequest & { stream?: false | undefined },
-  ): Promise<ProgressResponse>;
-  push(request: any): any {
-    return this.ollama.push(request);
-  }
-  abort(): void {
-    return this.ollama.abort();
-  }
-  encodeImage(image: string | Uint8Array): Promise<string> {
-    return this.ollama.encodeImage(image);
-  }
-  generate(
-    request: GenerateRequest & { stream: true },
-  ): Promise<AsyncGenerator<OllamaGenerateResponse>>;
-  generate(
-    request: GenerateRequest & { stream?: false | undefined },
-  ): Promise<OllamaGenerateResponse>;
-  generate(request: any): any {
-    return this.ollama.generate(request);
-  }
-  create(
-    request: CreateRequest & { stream: true },
-  ): Promise<AsyncGenerator<ProgressResponse>>;
-  create(
-    request: CreateRequest & { stream?: false | undefined },
-  ): Promise<ProgressResponse>;
-  create(request: any): any {
-    return this.ollama.create(request);
-  }
-  pull(
-    request: PullRequest & { stream: true },
-  ): Promise<AsyncGenerator<ProgressResponse>>;
-  pull(
-    request: PullRequest & { stream?: false | undefined },
-  ): Promise<ProgressResponse>;
-  pull(request: any): any {
-    return this.ollama.pull(request);
-  }
-  delete(request: DeleteRequest): Promise<StatusResponse> {
-    return this.ollama.delete(request);
-  }
-  copy(request: CopyRequest): Promise<StatusResponse> {
-    return this.ollama.copy(request);
-  }
-  list(): Promise<ListResponse> {
-    return this.ollama.list();
-  }
-  show(request: ShowRequest): Promise<ShowResponse> {
-    return this.ollama.show(request);
-  }
-  embeddings(request: EmbeddingsRequest): Promise<EmbeddingsResponse> {
-    return this.ollama.embeddings(request);
-  }
-}
+export { Ollama, type OllamaParams } from "@llamaindex/ollama";

packages/llamaindex/src/llm/types.ts

@@ -1,10 +0,0 @@
-import type { MessageContent } from "@llamaindex/core/llms";
-import type { NodeWithScore } from "@llamaindex/core/schema";
-
-export type RetrievalStartEvent = {
-  query: MessageContent;
-};
-export type RetrievalEndEvent = {
-  query: MessageContent;
-  nodes: NodeWithScore[];
-};

packages/llamaindex/src/objects/base.ts

@@ -1,8 +1,8 @@
 import type { BaseTool, MessageContent } from "@llamaindex/core/llms";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import type { BaseNode, Metadata } from "@llamaindex/core/schema";
 import { TextNode } from "@llamaindex/core/schema";
 import { extractText } from "@llamaindex/core/utils";
-import type { BaseRetriever } from "../Retriever.js";
 import type { VectorStoreIndex } from "../indices/vectorStore/index.js";
 
 // Assuming that necessary interfaces and classes (like OT, TextNode, BaseNode, etc.) are defined elsewhere
@@ -49,9 +49,6 @@ export abstract class BaseObjectNodeMapping {
 
 // You will need to implement specific subclasses of BaseObjectNodeMapping as per your project requirements.
 
-// todo: multimodal support
-type QueryType = MessageContent;
-
 export class ObjectRetriever<T = unknown> {
   _retriever: BaseRetriever;
   _objectNodeMapping: BaseObjectNodeMapping;
@@ -70,7 +67,7 @@ export class ObjectRetriever<T = unknown> {
   }
 
   // Translating the retrieve method
-  async retrieve(strOrQueryBundle: QueryType): Promise<T[]> {
+  async retrieve(strOrQueryBundle: MessageContent): Promise<T[]> {
     const nodes = await this.retriever.retrieve({
       query: extractText(strOrQueryBundle),
     });

packages/llamaindex/src/storage/docStore/KVDocumentStore.ts

@@ -29,7 +29,7 @@ export class KVDocumentStore extends BaseDocumentStore {
     for (const key in jsonDict) {
       const value = jsonDict[key];
       if (isValidDocJson(value)) {
-        docs[key] = jsonToDoc(value);
+        docs[key] = jsonToDoc(value, this.serializer);
       } else {
         console.warn(`Invalid JSON for docId ${key}`);
       }
@@ -52,7 +52,7 @@ export class KVDocumentStore extends BaseDocumentStore {
         );
       }
       const nodeKey = doc.id_;
-      const data = docToJson(doc);
+      const data = docToJson(doc, this.serializer);
       await this.kvstore.put(nodeKey, data, this.nodeCollection);
       const metadata: DocMetaData = { docHash: doc.hash };
 
@@ -94,7 +94,7 @@ export class KVDocumentStore extends BaseDocumentStore {
     if (!isValidDocJson(json)) {
       throw new Error(`Invalid JSON for docId ${docId}`);
     }
-    return jsonToDoc(json);
+    return jsonToDoc(json, this.serializer);
   }
 
   async getRefDocInfo(refDocId: string): Promise<RefDocInfo | undefined> {

packages/llamaindex/src/storage/docStore/PostgresDocumentStore.ts

@@ -1,19 +1,32 @@
 import { DEFAULT_NAMESPACE } from "@llamaindex/core/global";
-import { PostgresKVStore } from "../kvStore/PostgresKVStore.js";
+import {
+  PostgresKVStore,
+  type PostgresKVStoreConfig,
+} from "../kvStore/PostgresKVStore.js";
 import { KVDocumentStore } from "./KVDocumentStore.js";
+import { noneSerializer } from "./utils.js";
 
 const DEFAULT_TABLE_NAME = "llamaindex_doc_store";
 
+export type PostgresDocumentStoreConfig = PostgresKVStoreConfig & {
+  namespace?: string;
+};
+
 export class PostgresDocumentStore extends KVDocumentStore {
-  constructor(config?: {
-    schemaName?: string;
-    tableName?: string;
-    connectionString?: string;
-    namespace?: string;
-  }) {
+  serializer = noneSerializer;
+
+  constructor(config?: PostgresDocumentStoreConfig) {
     const kvStore = new PostgresKVStore({
       schemaName: config?.schemaName,
       tableName: config?.tableName || DEFAULT_TABLE_NAME,
+      ...(config && "clientConfig" in config
+        ? { clientConfig: config.clientConfig }
+        : config && "client" in config
+          ? {
+              client: config.client,
+              shouldConnect: config.shouldConnect ?? false,
+            }
+          : {}),
     });
     const namespace = config?.namespace || DEFAULT_NAMESPACE;
     super(kvStore, namespace);

packages/llamaindex/src/storage/docStore/types.ts

@@ -3,6 +3,7 @@ import {
   DEFAULT_PERSIST_DIR,
 } from "@llamaindex/core/global";
 import { BaseNode } from "@llamaindex/core/schema";
+import { jsonSerializer, type Serializer } from "./utils.js";
 
 const defaultPersistPath = `${DEFAULT_PERSIST_DIR}/${DEFAULT_DOC_STORE_PERSIST_FILENAME}`;
 
@@ -12,6 +13,8 @@ export interface RefDocInfo {
 }
 
 export abstract class BaseDocumentStore {
+  serializer: Serializer<any> = jsonSerializer;
+
   // Save/load
   persist(persistPath: string = defaultPersistPath): void {
     // Persist the docstore to a file.

packages/llamaindex/src/storage/docStore/utils.ts

@@ -4,12 +4,35 @@ import { Document, ObjectType, TextNode } from "@llamaindex/core/schema";
 const TYPE_KEY = "__type__";
 const DATA_KEY = "__data__";
 
-type DocJson = {
-  [TYPE_KEY]: ObjectType;
-  [DATA_KEY]: string;
+export interface Serializer<T> {
+  toPersistence(data: Record<string, unknown>): T;
+  fromPersistence(data: T): Record<string, unknown>;
+}
+
+export const jsonSerializer: Serializer<string> = {
+  toPersistence(data) {
+    return JSON.stringify(data);
+  },
+  fromPersistence(data) {
+    return JSON.parse(data);
+  },
 };
 
-export function isValidDocJson(docJson: any): docJson is DocJson {
+export const noneSerializer: Serializer<Record<string, unknown>> = {
+  toPersistence(data) {
+    return data;
+  },
+  fromPersistence(data) {
+    return data;
+  },
+};
+
+type DocJson<Data> = {
+  [TYPE_KEY]: ObjectType;
+  [DATA_KEY]: Data;
+};
+
+export function isValidDocJson(docJson: any): docJson is DocJson<unknown> {
   return (
     typeof docJson === "object" &&
     docJson !== null &&
@@ -18,16 +41,22 @@ export function isValidDocJson(docJson: any): docJson is DocJson {
   );
 }
 
-export function docToJson(doc: BaseNode): DocJson {
+export function docToJson(
+  doc: BaseNode,
+  serializer: Serializer<unknown>,
+): DocJson<unknown> {
   return {
-    [DATA_KEY]: JSON.stringify(doc.toJSON()),
+    [DATA_KEY]: serializer.toPersistence(doc.toJSON()),
     [TYPE_KEY]: doc.type,
   };
 }
 
-export function jsonToDoc(docDict: DocJson): BaseNode {
+export function jsonToDoc<Data>(
+  docDict: DocJson<Data>,
+  serializer: Serializer<Data>,
+): BaseNode {
   const docType = docDict[TYPE_KEY];
-  const dataDict = JSON.parse(docDict[DATA_KEY]);
+  const dataDict = serializer.fromPersistence(docDict[DATA_KEY]) as any;
   let doc: BaseNode;
 
   if (docType === ObjectType.DOCUMENT) {

packages/llamaindex/src/storage/indexStore/PostgresIndexStore.ts

@@ -1,19 +1,29 @@
 import { DEFAULT_NAMESPACE } from "@llamaindex/core/global";
-import { PostgresKVStore } from "../kvStore/PostgresKVStore.js";
+import {
+  PostgresKVStore,
+  type PostgresKVStoreConfig,
+} from "../kvStore/PostgresKVStore.js";
 import { KVIndexStore } from "./KVIndexStore.js";
 
 const DEFAULT_TABLE_NAME = "llamaindex_index_store";
 
+export type PostgresIndexStoreConfig = PostgresKVStoreConfig & {
+  namespace?: string;
+};
+
 export class PostgresIndexStore extends KVIndexStore {
-  constructor(config?: {
-    schemaName?: string;
-    tableName?: string;
-    connectionString?: string;
-    namespace?: string;
-  }) {
+  constructor(config?: PostgresIndexStoreConfig) {
     const kvStore = new PostgresKVStore({
       schemaName: config?.schemaName,
       tableName: config?.tableName || DEFAULT_TABLE_NAME,
+      ...(config && "clientConfig" in config
+        ? { clientConfig: config.clientConfig }
+        : config && "client" in config
+          ? {
+              client: config.client,
+              shouldConnect: config.shouldConnect ?? false,
+            }
+          : {}),
     });
     const namespace = config?.namespace || DEFAULT_NAMESPACE;
     super(kvStore, namespace);

packages/llamaindex/src/storage/kvStore/PostgresKVStore.ts

@@ -7,41 +7,76 @@ export type DataType = Record<string, Record<string, any>>;
 const DEFAULT_SCHEMA_NAME = "public";
 const DEFAULT_TABLE_NAME = "llamaindex_kv_store";
 
+export type PostgresKVStoreBaseConfig = {
+  schemaName?: string | undefined;
+  tableName?: string | undefined;
+};
+
+export type PostgresKVStoreClientConfig =
+  | {
+      /**
+       * Client configuration options for the pg client.
+       *
+       * {@link https://node-postgres.com/apis/client#new-client PostgresSQL Client API}
+       */
+      clientConfig?: pg.ClientConfig | undefined;
+    }
+  | {
+      /**
+       * A pg client or pool client instance.
+       * If provided, make sure it is not connected to the database yet, or it will throw an error.
+       */
+      shouldConnect?: boolean | undefined;
+      client?: pg.Client | pg.PoolClient;
+    };
+
+export type PostgresKVStoreConfig = PostgresKVStoreBaseConfig &
+  PostgresKVStoreClientConfig;
+
 export class PostgresKVStore extends BaseKVStore {
   private schemaName: string;
   private tableName: string;
-  private connectionString: string | undefined = undefined;
-  private db?: pg.Client;
 
-  constructor(config?: {
-    schemaName?: string | undefined;
-    tableName?: string | undefined;
-    connectionString?: string | undefined;
-  }) {
+  private isDBConnected: boolean = false;
+  private clientConfig: pg.ClientConfig | undefined = undefined;
+  private db?: pg.ClientBase | undefined = undefined;
+
+  constructor(config?: PostgresKVStoreConfig) {
     super();
     this.schemaName = config?.schemaName || DEFAULT_SCHEMA_NAME;
     this.tableName = config?.tableName || DEFAULT_TABLE_NAME;
-    this.connectionString = config?.connectionString;
-  }
-
-  private async getDb(): Promise<pg.Client> {
-    if (!this.db) {
-      try {
-        const pg = await import("pg");
-        const { Client } = pg.default ? pg.default : pg;
-        const db = new Client({ connectionString: this.connectionString });
-        await db.connect();
-        await this.checkSchema(db);
-        this.db = db;
-      } catch (err) {
-        console.error(err);
-        return Promise.reject(err instanceof Error ? err : new Error(`${err}`));
+    if (config) {
+      if ("clientConfig" in config) {
+        this.clientConfig = config.clientConfig;
+      } else if ("client" in config) {
+        this.isDBConnected =
+          config?.shouldConnect !== undefined ? !config.shouldConnect : false;
+        this.db = config.client;
       }
     }
-    return Promise.resolve(this.db);
   }
 
-  private async checkSchema(db: pg.Client) {
+  private async getDb(): Promise<pg.ClientBase> {
+    if (!this.db) {
+      const pg = await import("pg");
+      const { Client } = pg.default ? pg.default : pg;
+      const db = new Client({ ...this.clientConfig });
+      await db.connect();
+      this.isDBConnected = true;
+      this.db = db;
+    }
+    if (this.db && !this.isDBConnected) {
+      await this.db.connect();
+      this.isDBConnected = true;
+    }
+    this.db.on("end", () => {
+      this.isDBConnected = false;
+    });
+    await this.checkSchema(this.db);
+    return this.db;
+  }
+
+  private async checkSchema(db: pg.ClientBase) {
     await db.query(`CREATE SCHEMA IF NOT EXISTS ${this.schemaName}`);
     const tbl = `CREATE TABLE IF NOT EXISTS ${this.schemaName}.${this.tableName} (
       id uuid DEFAULT gen_random_uuid() PRIMARY KEY,
@@ -97,7 +132,7 @@ export class PostgresKVStore extends BaseKVStore {
       const sql = `SELECT * FROM ${this.schemaName}.${this.tableName} WHERE key = $1 AND collection = $2`;
       const result = await db.query(sql, [key, collection]);
       await db.query("COMMIT");
-      return result.rows[0].value;
+      return result.rows[0]?.value;
     } catch (error) {
       await db.query("ROLLBACK");
       throw error;

packages/llamaindex/src/vector-store/PGVectorStore.ts

@@ -14,25 +14,44 @@ import {
 import { escapeLikeString } from "./utils.js";
 
 import type { BaseEmbedding } from "@llamaindex/core/embeddings";
+import { DEFAULT_COLLECTION } from "@llamaindex/core/global";
 import type { BaseNode, Metadata } from "@llamaindex/core/schema";
 import { Document, MetadataMode } from "@llamaindex/core/schema";
 
 export const PGVECTOR_SCHEMA = "public";
 export const PGVECTOR_TABLE = "llamaindex_embedding";
+export const DEFAULT_DIMENSIONS = 1536;
 
-export type PGVectorStoreConfig = Pick<
-  pg.ClientConfig,
-  "user" | "database" | "password" | "connectionString"
-> & {
+type PGVectorStoreBaseConfig = {
   schemaName?: string | undefined;
   tableName?: string | undefined;
   dimensions?: number | undefined;
   embedModel?: BaseEmbedding | undefined;
 };
 
+export type PGVectorStoreConfig = PGVectorStoreBaseConfig &
+  (
+    | {
+        /**
+         * Client configuration options for the pg client.
+         *
+         * {@link https://node-postgres.com/apis/client#new-client PostgresSQL Client API}
+         */
+        clientConfig: pg.ClientConfig;
+      }
+    | {
+        /**
+         * A pg client or pool client instance.
+         * If provided, make sure it is not connected to the database yet, or it will throw an error.
+         */
+        shouldConnect?: boolean | undefined;
+        client: pg.Client | pg.PoolClient;
+      }
+  );
+
 /**
  * Provides support for writing and querying vector data in Postgres.
- * Note: Can't be used with data created using the Python version of the vector store (https://docs.llamaindex.ai/en/stable/examples/vector_stores/postgres.html)
+ * Note: Can't be used with data created using the Python version of the vector store (https://docs.llamaindex.ai/en/stable/examples/vector_stores/postgres/)
  */
 export class PGVectorStore
   extends VectorStoreBase
@@ -40,52 +59,26 @@ export class PGVectorStore
 {
   storesText: boolean = true;
 
-  private collection: string = "";
-  private schemaName: string = PGVECTOR_SCHEMA;
-  private tableName: string = PGVECTOR_TABLE;
+  private collection: string = DEFAULT_COLLECTION;
+  private readonly schemaName: string = PGVECTOR_SCHEMA;
+  private readonly tableName: string = PGVECTOR_TABLE;
+  private readonly dimensions: number = DEFAULT_DIMENSIONS;
 
-  private user: pg.ClientConfig["user"] | undefined = undefined;
-  private password: pg.ClientConfig["password"] | undefined = undefined;
-  private database: pg.ClientConfig["database"] | undefined = undefined;
-  private connectionString: pg.ClientConfig["connectionString"] | undefined =
-    undefined;
+  private isDBConnected: boolean = false;
+  private db: pg.ClientBase | null = null;
+  private readonly clientConfig: pg.ClientConfig | null = null;
 
-  private dimensions: number = 1536;
-
-  private db?: pg.ClientBase;
-
-  /**
-   * Constructs a new instance of the PGVectorStore
-   *
-   * If the `connectionString` is not provided the following env variables are
-   * used to connect to the DB:
-   * PGHOST=your database host
-   * PGUSER=your database user
-   * PGPASSWORD=your database password
-   * PGDATABASE=your database name
-   * PGPORT=your database port
-   */
-  constructor(configOrClient?: PGVectorStoreConfig | pg.ClientBase) {
-    // We cannot import pg from top level, it might have side effects
-    //  so we only check if the config.connect function exists
-    if (
-      configOrClient &&
-      "connect" in configOrClient &&
-      typeof configOrClient.connect === "function"
-    ) {
-      const db = configOrClient as pg.ClientBase;
-      super();
-      this.db = db;
+  constructor(config: PGVectorStoreConfig) {
+    super(config?.embedModel);
+    this.schemaName = config?.schemaName ?? PGVECTOR_SCHEMA;
+    this.tableName = config?.tableName ?? PGVECTOR_TABLE;
+    this.dimensions = config?.dimensions ?? DEFAULT_DIMENSIONS;
+    if ("clientConfig" in config) {
+      this.clientConfig = config.clientConfig;
     } else {
-      const config = configOrClient as PGVectorStoreConfig;
-      super(config?.embedModel);
-      this.schemaName = config?.schemaName ?? PGVECTOR_SCHEMA;
-      this.tableName = config?.tableName ?? PGVECTOR_TABLE;
-      this.user = config?.user;
-      this.password = config?.password;
-      this.database = config?.database;
-      this.connectionString = config?.connectionString;
-      this.dimensions = config?.dimensions ?? 1536;
+      this.isDBConnected =
+        config.shouldConnect !== undefined ? !config.shouldConnect : false;
+      this.db = config.client;
     }
   }
 
@@ -113,39 +106,41 @@ export class PGVectorStore
 
   private async getDb(): Promise<pg.ClientBase> {
     if (!this.db) {
-      try {
-        const pg = await import("pg");
-        const { Client } = pg.default ? pg.default : pg;
+      const pg = await import("pg");
+      const { Client } = pg.default ? pg.default : pg;
 
-        const { registerType } = await import("pgvector/pg");
-        // Create DB connection
-        // Read connection params from env - see comment block above
-        const db = new Client({
-          user: this.user,
-          password: this.password,
-          database: this.database,
-          connectionString: this.connectionString,
-        });
-        await db.connect();
+      const { registerTypes } = await import("pgvector/pg");
+      // Create DB connection
+      // Read connection params from env - see comment block above
+      const db = new Client({
+        ...this.clientConfig,
+      });
 
-        // Check vector extension
-        await db.query("CREATE EXTENSION IF NOT EXISTS vector");
-        await registerType(db);
+      await db.connect();
+      this.isDBConnected = true;
 
-        // All good?  Keep the connection reference
-        this.db = db;
-      } catch (err) {
-        console.error(err);
-        return Promise.reject(err instanceof Error ? err : new Error(`${err}`));
-      }
+      // Check vector extension
+      await db.query("CREATE EXTENSION IF NOT EXISTS vector");
+      await registerTypes(db);
+
+      // All good?  Keep the connection reference
+      this.db = db;
     }
 
-    const db = this.db;
+    if (this.db && !this.isDBConnected) {
+      await this.db.connect();
+      this.isDBConnected = true;
+    }
+
+    this.db.on("end", () => {
+      // Connection closed
+      this.isDBConnected = false;
+    });
 
     // Check schema, table(s), index(es)
-    await this.checkSchema(db);
+    await this.checkSchema(this.db);
 
-    return Promise.resolve(this.db);
+    return this.db;
   }
 
   private async checkSchema(db: pg.ClientBase) {

packages/llamaindex/src/vector-store/UpstashVectorStore.ts

@@ -0,0 +1,237 @@
+import {
+  VectorStoreBase,
+  type IEmbedModel,
+  type MetadataFilter,
+  type MetadataFilters,
+  type VectorStoreNoEmbedModel,
+  type VectorStoreQuery,
+  type VectorStoreQueryResult,
+} from "./types.js";
+
+import type { BaseNode, Metadata, TextNode } from "@llamaindex/core/schema";
+import { getEnv } from "@llamaindex/env";
+import { Index } from "@upstash/vector";
+import { metadataDictToNode, nodeToMetadata } from "./utils.js";
+
+type UpstashParams = {
+  namespace?: string;
+  token?: string;
+  endpoint?: string;
+  maxBatchSize?: number;
+} & IEmbedModel;
+
+/**
+ * Provides support for writing and querying vector data in Upstash.
+ */
+export class UpstashVectorStore
+  extends VectorStoreBase
+  implements VectorStoreNoEmbedModel
+{
+  storesText: boolean = true;
+
+  private db: Index;
+  private maxBatchSize: number;
+  namespace: string;
+
+  /**
+   * @param namespace namespace to use
+   * @param token upstash vector token. if not set, `process.env.UPSTASH_VECTOR_REST_TOKEN` is used.
+   * @param endpoint upstash vector endpoint. If not set, `process.env.UPSTASH_VECTOR_REST_URL` is used.
+   * @param maxBatchSize maximum number of vectors upserted at once. Default is 1000.
+   *
+   * @example
+   * ```ts
+   * const vectorStore = new UpstashVectorStore({ namespace: "my-namespace" })
+   * ```
+   */
+  constructor(params?: UpstashParams) {
+    super(params?.embedModel);
+    this.namespace = params?.namespace ?? "";
+    this.maxBatchSize = params?.maxBatchSize ?? 1000;
+    const token = params?.token ?? getEnv("UPSTASH_VECTOR_REST_TOKEN");
+    const endpoint = params?.endpoint ?? getEnv("UPSTASH_VECTOR_REST_URL");
+
+    if (!token) {
+      throw new Error(
+        "Must specify UPSTASH_VECTOR_REST_TOKEN via env variable.",
+      );
+    }
+    if (!endpoint) {
+      throw new Error("Must specify UPSTASH_VECTOR_REST_URL via env variable.");
+    }
+    this.db = new Index({ token, url: endpoint });
+  }
+
+  private async getDb(): Promise<Index> {
+    if (!this.db) {
+      const { Index } = await import("@upstash/vector");
+      this.db = new Index();
+    }
+
+    return this.db;
+  }
+
+  /**
+   * Connects to the database specified in environment vars.
+   * @returns A connection to the database, or the error encountered while connecting/setting up.
+   */
+  client(): Promise<Index> {
+    return this.getDb();
+  }
+
+  /**
+   * Adds vector record(s) to the table.
+   * @param embeddingResults The Nodes to be inserted, optionally including metadata tuples.
+   * @returns ids of the embeddings (infered from the id_ field of embeddingResults objects)
+   */
+  async add(embeddingResults: BaseNode<Metadata>[]): Promise<string[]> {
+    if (embeddingResults.length == 0) {
+      return [];
+    }
+
+    const nodes = embeddingResults.map(this.nodeToRecord);
+    const result = await this.upsertInBatches(nodes);
+    if (result != "OK") {
+      throw new Error("Failed to save chunk");
+    }
+    return nodes.map((node) => node.id);
+  }
+
+  /**
+   * Adds plain text record(s) to the table. Upstash take cares of embedding conversion.
+   * @param text The Nodes to be inserted, optionally including metadata tuples.
+   * @returns ids of the embeddings (infered from the id_ field of embeddingResults objects)
+   */
+  async addPlainText(text: TextNode<Metadata>[]): Promise<string[]> {
+    if (text.length == 0) {
+      return [];
+    }
+
+    const nodes = text.map(this.textNodeToRecord);
+    const result = await this.upsertInBatches(nodes);
+    if (result != "OK") {
+      throw new Error("Failed to save chunk");
+    }
+    return nodes.map((node) => node.id);
+  }
+
+  private async upsertInBatches(
+    nodes:
+      | ReturnType<UpstashVectorStore["textNodeToRecord"]>[]
+      | ReturnType<UpstashVectorStore["nodeToRecord"]>[],
+  ) {
+    const promises: Promise<string>[] = [];
+    for (let i = 0; i < nodes.length; i += this.maxBatchSize) {
+      const batch = nodes.slice(i, i + this.maxBatchSize);
+      promises.push(this.db.upsert(batch, { namespace: this.namespace }));
+    }
+    const results = await Promise.all(promises);
+    return results.every((result) => result === "OK") ? "OK" : "NOT-OK";
+  }
+
+  /**
+   * Deletes a single record from the database by id.
+   * NOTE: Uses the collection property controlled by setCollection/getCollection.
+   * @param refDocId Unique identifier for the record to delete.
+   * @returns Promise that resolves if the delete query did not throw an error.
+   */
+  async delete(refDocId: string): Promise<void> {
+    await this.db.namespace(this.namespace).delete(refDocId);
+  }
+
+  /**
+   * Deletes a single record from the database by id.
+   * NOTE: Uses the collection property controlled by setCollection/getCollection.
+   * @param refDocId Unique identifier for the record to delete.
+   * @param deleteKwargs Required by VectorStore interface.  Currently ignored.
+   * @returns Promise that resolves if the delete query did not throw an error.
+   */
+  async deleteMany(refDocId: string[]): Promise<void> {
+    await this.db.namespace(this.namespace).delete(refDocId);
+  }
+
+  /**
+   * Query the vector store for the closest matching data to the query embeddings
+   * @param query The VectorStoreQuery to be used
+   * @param options Required by VectorStore interface.  Currently ignored.
+   * @returns Zero or more Document instances with data from the vector store.
+   */
+  async query(
+    query: VectorStoreQuery,
+    _options?: any,
+  ): Promise<VectorStoreQueryResult> {
+    const filter = this.toUpstashFilter(query.filters);
+
+    const defaultOptions: any = {
+      vector: query.queryEmbedding,
+      topK: query.similarityTopK,
+      includeVectors: true,
+      includeMetadata: true,
+      filter,
+    };
+
+    const db = this.db;
+    const results = await db.query(defaultOptions, {
+      namespace: this.namespace,
+    });
+
+    const nodes = results.map((result) => {
+      const node = metadataDictToNode(result.metadata as Record<string, any>, {
+        fallback: {
+          id: result.id,
+          metadata: result.metadata,
+          embedding: result.vector,
+        },
+      });
+      return node;
+    });
+
+    const ret = {
+      nodes: nodes,
+      similarities: results.map((row) => row.score || 999),
+      ids: results.map((row) => String(row.id)),
+    };
+
+    return ret;
+  }
+
+  toFilterString(filter: MetadataFilter) {
+    return `${filter.key} ${filter.operator} ${filter.value}`;
+  }
+
+  toUpstashFilter(stdFilters?: MetadataFilters) {
+    if (!stdFilters?.filters) return;
+
+    for (const item of stdFilters.filters) {
+      if (item.operator === "==") {
+        //@ts-expect-error Upstash equal operator uses only one equal sign, so we have to replace it.
+        item.operator = "=";
+      }
+    }
+
+    const filterStrings = stdFilters.filters.map(this.toFilterString);
+
+    if (filterStrings.length === 1) {
+      return filterStrings[0];
+    }
+    return filterStrings.join(` ${stdFilters.condition ?? "and"} `);
+  }
+
+  nodeToRecord(node: BaseNode<Metadata>) {
+    const id: any = node.id_.length ? node.id_ : null;
+    return {
+      id: id,
+      vector: node.getEmbedding(),
+      metadata: nodeToMetadata(node),
+    };
+  }
+
+  textNodeToRecord(node: TextNode<Metadata>) {
+    const id: any = node.id_.length ? node.id_ : null;
+    return {
+      id,
+      data: node.text,
+      metadata: nodeToMetadata(node),
+    };
+  }
+}

packages/llm/groq/CHANGELOG.md

@@ -1,5 +1,17 @@
 # @llamaindex/groq
 
+## 0.0.7
+
+### Patch Changes
+
+- @llamaindex/openai@0.1.8
+
+## 0.0.6
+
+### Patch Changes
+
+- @llamaindex/openai@0.1.7
+
 ## 0.0.5
 
 ### Patch Changes

packages/llm/groq/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/groq",
   "description": "Groq Adapter for LlamaIndex",
-  "version": "0.0.5",
+  "version": "0.0.7",
   "type": "module",
   "main": "./dist/index.cjs",
   "module": "./dist/index.js",

packages/llm/ollama/package.json

@@ -0,0 +1,41 @@
+{
+  "name": "@llamaindex/ollama",
+  "description": "Ollama Adapter for LlamaIndex",
+  "version": "0.0.1",
+  "type": "module",
+  "main": "./dist/index.cjs",
+  "module": "./dist/index.js",
+  "exports": {
+    ".": {
+      "require": {
+        "types": "./dist/index.d.cts",
+        "default": "./dist/index.cjs"
+      },
+      "import": {
+        "types": "./dist/index.d.ts",
+        "default": "./dist/index.js"
+      }
+    }
+  },
+  "files": [
+    "dist"
+  ],
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/run-llama/LlamaIndexTS.git",
+    "directory": "packages/llm/openai"
+  },
+  "scripts": {
+    "build": "bunchee",
+    "dev": "bunchee --watch"
+  },
+  "devDependencies": {
+    "bunchee": "5.3.2"
+  },
+  "dependencies": {
+    "@llamaindex/core": "workspace:*",
+    "@llamaindex/env": "workspace:*",
+    "ollama": "^0.5.9",
+    "remeda": "^2.12.0"
+  }
+}

packages/llm/ollama/src/index.ts

@@ -0,0 +1,172 @@
+import { BaseEmbedding } from "@llamaindex/core/embeddings";
+import type {
+  ChatResponse,
+  ChatResponseChunk,
+  CompletionResponse,
+  LLM,
+  LLMChatParamsNonStreaming,
+  LLMChatParamsStreaming,
+  LLMCompletionParamsNonStreaming,
+  LLMCompletionParamsStreaming,
+  LLMMetadata,
+} from "@llamaindex/core/llms";
+import { extractText, streamConverter } from "@llamaindex/core/utils";
+import {
+  Ollama as OllamaBase,
+  type Config,
+  type ChatResponse as OllamaChatResponse,
+  type GenerateResponse as OllamaGenerateResponse,
+  type Options,
+} from "ollama/browser";
+
+const messageAccessor = (part: OllamaChatResponse): ChatResponseChunk => {
+  return {
+    raw: part,
+    delta: part.message.content,
+  };
+};
+
+const completionAccessor = (
+  part: OllamaGenerateResponse,
+): CompletionResponse => {
+  return { text: part.response, raw: part };
+};
+
+export type OllamaParams = {
+  model: string;
+  config?: Partial<Config>;
+  options?: Partial<Options>;
+};
+
+export class Ollama extends BaseEmbedding implements LLM {
+  public readonly ollama: OllamaBase;
+
+  // https://ollama.ai/library
+  model: string;
+
+  options: Partial<Omit<Options, "num_ctx" | "top_p" | "temperature">> &
+    Pick<Options, "num_ctx" | "top_p" | "temperature"> = {
+    num_ctx: 4096,
+    top_p: 0.9,
+    temperature: 0.7,
+  };
+
+  constructor(params: OllamaParams) {
+    super();
+    this.model = params.model;
+    this.ollama = new OllamaBase(params.config);
+    if (params.options) {
+      this.options = {
+        ...this.options,
+        ...params.options,
+      };
+    }
+  }
+
+  get metadata(): LLMMetadata {
+    const { temperature, top_p, num_ctx } = this.options;
+    return {
+      model: this.model,
+      temperature: temperature,
+      topP: top_p,
+      maxTokens: this.options.num_ctx,
+      contextWindow: num_ctx,
+      tokenizer: undefined,
+    };
+  }
+
+  chat(
+    params: LLMChatParamsStreaming,
+  ): Promise<AsyncIterable<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    const { messages, stream } = params;
+    const payload = {
+      model: this.model,
+      messages: messages.map((message) => ({
+        role: message.role,
+        content: extractText(message.content),
+      })),
+      stream: !!stream,
+      options: {
+        ...this.options,
+      },
+    };
+    if (!stream) {
+      const chatResponse = await this.ollama.chat({
+        ...payload,
+        stream: false,
+      });
+
+      return {
+        message: {
+          role: "assistant",
+          content: chatResponse.message.content,
+        },
+        raw: chatResponse,
+      };
+    } else {
+      const stream = await this.ollama.chat({
+        ...payload,
+        stream: true,
+      });
+      return streamConverter(stream, messageAccessor);
+    }
+  }
+
+  complete(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<CompletionResponse>>;
+  complete(
+    params: LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse>;
+  async complete(
+    params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
+    const { prompt, stream } = params;
+    const payload = {
+      model: this.model,
+      prompt: extractText(prompt),
+      stream: !!stream,
+      options: {
+        ...this.options,
+      },
+    };
+    if (!stream) {
+      const response = await this.ollama.generate({
+        ...payload,
+        stream: false,
+      });
+      return {
+        text: response.response,
+        raw: response,
+      };
+    } else {
+      const stream = await this.ollama.generate({
+        ...payload,
+        stream: true,
+      });
+      return streamConverter(stream, completionAccessor);
+    }
+  }
+
+  private async getEmbedding(prompt: string): Promise<number[]> {
+    const payload = {
+      model: this.model,
+      prompt,
+      options: {
+        ...this.options,
+      },
+    };
+    const response = await this.ollama.embeddings({
+      ...payload,
+    });
+    return response.embedding;
+  }
+
+  async getTextEmbedding(text: string): Promise<number[]> {
+    return this.getEmbedding(text);
+  }
+}

packages/llm/ollama/tsconfig.json

@@ -0,0 +1,18 @@
+{
+  "extends": "../../../tsconfig.json",
+  "compilerOptions": {
+    "target": "ESNext",
+    "module": "ESNext",
+    "moduleResolution": "bundler",
+    "outDir": "./lib"
+  },
+  "include": ["./src"],
+  "references": [
+    {
+      "path": "../../llamaindex/tsconfig.json"
+    },
+    {
+      "path": "../../env/tsconfig.json"
+    }
+  ]
+}

packages/llm/openai/CHANGELOG.md

@@ -1,5 +1,19 @@
 # @llamaindex/openai
 
+## 0.1.8
+
+### Patch Changes
+
+- Updated dependencies [8b7fdba]
+  - @llamaindex/core@0.2.6
+
+## 0.1.7
+
+### Patch Changes
+
+- Updated dependencies [d902cc3]
+  - @llamaindex/core@0.2.5
+
 ## 0.1.6
 
 ### Patch Changes

packages/llm/openai/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/openai",
   "description": "OpenAI Adapter for LlamaIndex",
-  "version": "0.1.6",
+  "version": "0.1.8",
   "type": "module",
   "main": "./dist/index.cjs",
   "module": "./dist/index.js",

tsconfig.json

@@ -32,6 +32,9 @@
     {
       "path": "./packages/llm/groq/tsconfig.json"
     },
+    {
+      "path": "./packages/llm/ollama/tsconfig.json"
+    },
     {
       "path": "./packages/cloud/tsconfig.json"
     },