mirror of
https://github.com/run-llama/LlamaIndexTS.git
synced 2026-07-04 03:40:26 -04:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bdc4bfe7b0 | |||
| 025ffe6b50 | |||
| a6595747fa | |||
| d902cc3e7e | |||
| 726eb41359 | |||
| e9714dbfcd | |||
| a3618e761e |
@@ -1,5 +1,21 @@
|
||||
# docs
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "docs",
|
||||
"version": "0.0.73",
|
||||
"version": "0.0.75",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
"docusaurus": "docusaurus",
|
||||
|
||||
@@ -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();
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
|
||||
@@ -1,5 +1,21 @@
|
||||
# @llamaindex/autotool
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,5 +1,23 @@
|
||||
# @llamaindex/autotool-01-node-example
|
||||
|
||||
## 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
|
||||
|
||||
@@ -13,5 +13,5 @@
|
||||
"scripts": {
|
||||
"start": "node --import tsx --import @llamaindex/autotool/node ./src/index.ts"
|
||||
},
|
||||
"version": "0.0.13"
|
||||
"version": "0.0.15"
|
||||
}
|
||||
|
||||
@@ -1,5 +1,23 @@
|
||||
# @llamaindex/autotool-02-next-example
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "@llamaindex/autotool-02-next-example",
|
||||
"private": true,
|
||||
"version": "0.1.57",
|
||||
"version": "0.1.59",
|
||||
"scripts": {
|
||||
"dev": "next dev",
|
||||
"build": "next build",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "@llamaindex/autotool",
|
||||
"type": "module",
|
||||
"version": "3.0.4",
|
||||
"version": "3.0.6",
|
||||
"description": "auto transpile your JS function to LLM Agent compatible",
|
||||
"files": [
|
||||
"dist",
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"extends": ["//"],
|
||||
"tasks": {
|
||||
"build": {
|
||||
"outputs": ["dist/**", "src/client/**"]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,12 @@
|
||||
# @llamaindex/community
|
||||
|
||||
## 0.0.39
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [d902cc3]
|
||||
- @llamaindex/core@0.2.5
|
||||
|
||||
## 0.0.38
|
||||
|
||||
### Patch Changes
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "@llamaindex/community",
|
||||
"description": "Community package for LlamaIndexTS",
|
||||
"version": "0.0.38",
|
||||
"version": "0.0.39",
|
||||
"type": "module",
|
||||
"types": "dist/type/index.d.ts",
|
||||
"main": "dist/cjs/index.js",
|
||||
|
||||
@@ -1,5 +1,11 @@
|
||||
# @llamaindex/core
|
||||
|
||||
## 0.2.5
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- d902cc3: Fix context not being sent using ContextChatEngine
|
||||
|
||||
## 0.2.4
|
||||
|
||||
### Patch Changes
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "@llamaindex/core",
|
||||
"type": "module",
|
||||
"version": "0.2.4",
|
||||
"version": "0.2.5",
|
||||
"description": "LlamaIndex Core Module",
|
||||
"exports": {
|
||||
"./node-parser": {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -1,5 +1,21 @@
|
||||
# @llamaindex/experimental
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "@llamaindex/experimental",
|
||||
"description": "Experimental package for LlamaIndexTS",
|
||||
"version": "0.0.82",
|
||||
"version": "0.0.84",
|
||||
"type": "module",
|
||||
"types": "dist/type/index.d.ts",
|
||||
"main": "dist/cjs/index.js",
|
||||
|
||||
@@ -1,5 +1,27 @@
|
||||
# llamaindex
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,5 +1,21 @@
|
||||
# @llamaindex/cloudflare-worker-agent-test
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@llamaindex/cloudflare-worker-agent-test",
|
||||
"version": "0.0.66",
|
||||
"version": "0.0.68",
|
||||
"type": "module",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
|
||||
@@ -1,5 +1,21 @@
|
||||
# @llamaindex/next-agent-test
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@llamaindex/next-agent-test",
|
||||
"version": "0.1.66",
|
||||
"version": "0.1.68",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
"dev": "next dev",
|
||||
|
||||
@@ -1,5 +1,21 @@
|
||||
# test-edge-runtime
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@llamaindex/nextjs-edge-runtime-test",
|
||||
"version": "0.1.65",
|
||||
"version": "0.1.67",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
"dev": "next dev",
|
||||
|
||||
@@ -1,5 +1,21 @@
|
||||
# @llamaindex/next-node-runtime
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@llamaindex/next-node-runtime-test",
|
||||
"version": "0.0.47",
|
||||
"version": "0.0.49",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
"dev": "next dev",
|
||||
|
||||
@@ -1,5 +1,21 @@
|
||||
# @llamaindex/waku-query-engine-test
|
||||
|
||||
## 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
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@llamaindex/waku-query-engine-test",
|
||||
"version": "0.0.66",
|
||||
"version": "0.0.68",
|
||||
"type": "module",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
|
||||
@@ -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;
|
||||
});
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "llamaindex",
|
||||
"version": "0.6.4",
|
||||
"version": "0.6.6",
|
||||
"license": "MIT",
|
||||
"type": "module",
|
||||
"keywords": [
|
||||
@@ -43,6 +43,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",
|
||||
|
||||
@@ -356,9 +356,8 @@ export abstract class AgentRunner<
|
||||
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>[];
|
||||
|
||||
@@ -78,9 +78,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({
|
||||
@@ -103,7 +101,7 @@ export class CondenseQuestionChatEngine
|
||||
? new ChatMemoryBuffer({
|
||||
chatHistory:
|
||||
params.chatHistory instanceof BaseMemory
|
||||
? await params.chatHistory.getMessages(message)
|
||||
? await params.chatHistory.getMessages()
|
||||
: params.chatHistory,
|
||||
})
|
||||
: this.chatHistory;
|
||||
|
||||
@@ -92,7 +92,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
|
||||
? new ChatMemoryBuffer({
|
||||
chatHistory:
|
||||
params.chatHistory instanceof BaseMemory
|
||||
? await params.chatHistory.getMessages(message)
|
||||
? await params.chatHistory.getMessages()
|
||||
: params.chatHistory,
|
||||
})
|
||||
: this.chatHistory;
|
||||
@@ -139,7 +139,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 };
|
||||
}
|
||||
|
||||
|
||||
@@ -40,7 +40,7 @@ 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;
|
||||
@@ -48,7 +48,7 @@ export class SimpleChatEngine implements ChatEngine {
|
||||
|
||||
if (stream) {
|
||||
const stream = await this.llm.chat({
|
||||
messages: await chatHistory.getMessages(params.message),
|
||||
messages: await chatHistory.getMessages(),
|
||||
stream: true,
|
||||
});
|
||||
return streamConverter(
|
||||
@@ -66,7 +66,7 @@ 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);
|
||||
|
||||
@@ -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 { KVDocumentStore } from "./KVDocumentStore.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;
|
||||
}) {
|
||||
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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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),
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,11 @@
|
||||
# @llamaindex/groq
|
||||
|
||||
## 0.0.6
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- @llamaindex/openai@0.1.7
|
||||
|
||||
## 0.0.5
|
||||
|
||||
### Patch Changes
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "@llamaindex/groq",
|
||||
"description": "Groq Adapter for LlamaIndex",
|
||||
"version": "0.0.5",
|
||||
"version": "0.0.6",
|
||||
"type": "module",
|
||||
"main": "./dist/index.cjs",
|
||||
"module": "./dist/index.js",
|
||||
|
||||
@@ -1,5 +1,12 @@
|
||||
# @llamaindex/openai
|
||||
|
||||
## 0.1.7
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [d902cc3]
|
||||
- @llamaindex/core@0.2.5
|
||||
|
||||
## 0.1.6
|
||||
|
||||
### Patch Changes
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "@llamaindex/openai",
|
||||
"description": "OpenAI Adapter for LlamaIndex",
|
||||
"version": "0.1.6",
|
||||
"version": "0.1.7",
|
||||
"type": "module",
|
||||
"main": "./dist/index.cjs",
|
||||
"module": "./dist/index.js",
|
||||
|
||||
Generated
+8
@@ -565,6 +565,9 @@ importers:
|
||||
'@types/pg':
|
||||
specifier: ^8.11.8
|
||||
version: 8.11.8
|
||||
'@upstash/vector':
|
||||
specifier: ^1.1.5
|
||||
version: 1.1.5
|
||||
'@zilliz/milvus2-sdk-node':
|
||||
specifier: ^2.4.6
|
||||
version: 2.4.6
|
||||
@@ -4511,6 +4514,9 @@ packages:
|
||||
'@ungap/structured-clone@1.2.0':
|
||||
resolution: {integrity: sha512-zuVdFrMJiuCDQUMCzQaD6KL28MjnqqN8XnAqiEq9PNm/hCPTSGfrXCOfwj1ow4LFb/tNymJPwsNbVePc1xFqrQ==}
|
||||
|
||||
'@upstash/vector@1.1.5':
|
||||
resolution: {integrity: sha512-55+Beu/kCwjcnzg6fnMN06v9PYU1lv9NQfQwpjrJAQTH8GOprcRsQeyXBdNHKNzoQvRnVS0ENd5CDgFoljfrAw==}
|
||||
|
||||
'@vitejs/plugin-react@4.3.1':
|
||||
resolution: {integrity: sha512-m/V2syj5CuVnaxcUJOQRel/Wr31FFXRFlnOoq1TVtkCxsY5veGMTEmpWHndrhB2U8ScHtCQB1e+4hWYExQc6Lg==}
|
||||
engines: {node: ^14.18.0 || >=16.0.0}
|
||||
@@ -16426,6 +16432,8 @@ snapshots:
|
||||
|
||||
'@ungap/structured-clone@1.2.0': {}
|
||||
|
||||
'@upstash/vector@1.1.5': {}
|
||||
|
||||
'@vitejs/plugin-react@4.3.1(vite@5.4.2(@types/node@22.5.5)(terser@5.32.0))':
|
||||
dependencies:
|
||||
'@babel/core': 7.25.2
|
||||
|
||||
Reference in New Issue
Block a user