mirror of
https://github.com/run-llama/LlamaIndexTS.git
synced 2026-07-15 06:52:45 -04:00
Compare commits
296 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 3bab23172a | |||
| 18c132d494 | |||
| d072353e08 | |||
| c8bbc101cc | |||
| b93f748998 | |||
| ecb100448a | |||
| c749c856b5 | |||
| 0baf278972 | |||
| ae7780266a | |||
| 587960aebe | |||
| 4e1b6784f7 | |||
| 8b381f2640 | |||
| 0dc7fa6c34 | |||
| 2a2bf682bf | |||
| 87526129fb | |||
| 8ed1b7aa46 | |||
| 4084bd0ecc | |||
| d11eaceaf1 | |||
| 1e6986fbc5 | |||
| 11a19bdec7 | |||
| 51064f1b90 | |||
| 3385cd19e8 | |||
| 852f8517df | |||
| bb917f9818 | |||
| 10248fb29f | |||
| 446dc85bdd | |||
| 4aa2c226a9 | |||
| bf9ba8313a | |||
| 444b59c557 | |||
| b2e1df94db | |||
| b4963cabc8 | |||
| 2851024340 | |||
| 7f25a25729 | |||
| acfe23265a | |||
| 2c6fbbd7dd | |||
| f84507f513 | |||
| be6a9e4a48 | |||
| 69e7634619 | |||
| d18748aba4 | |||
| 27c4ef3410 | |||
| a7ee392d3e | |||
| 4415a6fdef | |||
| 1e1e6e96a1 | |||
| 461d1dfbcc | |||
| 3e8c923641 | |||
| 5975fafefb | |||
| 71169fd545 | |||
| be895d564d | |||
| f36a27c218 | |||
| 8cdb07f151 | |||
| ea403a0ffe | |||
| 7f0b4e66ae | |||
| 3b226965ba | |||
| 63daf77412 | |||
| df5cbe30a6 | |||
| 9e1a536778 | |||
| a1db8833ef | |||
| 95dd0e0158 | |||
| 079a1d5cc3 | |||
| 2377d1a466 | |||
| 9f9f29391e | |||
| b64716d3f7 | |||
| d7a47abe38 | |||
| 58b314a61e | |||
| 4431ec7a5e | |||
| 9542026d70 | |||
| cc4c5b64c0 | |||
| 82c2aac4a0 | |||
| a143e0f0f1 | |||
| db9775dc32 | |||
| 538c0b0740 | |||
| 21cd88caf6 | |||
| 0660d9e2a5 | |||
| 25257f49d7 | |||
| dd615f106d | |||
| 5db64d61e0 | |||
| ee5e1f94e4 | |||
| 031e926414 | |||
| 88b4b3143d | |||
| c1ce84ecec | |||
| d670011363 | |||
| c88332366b | |||
| cfee282c28 | |||
| 91b42a3539 | |||
| 02b1d176c5 | |||
| 63d072b8cc | |||
| 256d44f255 | |||
| e2a6805a31 | |||
| d46fc12079 | |||
| 5ce88f107c | |||
| 683c4addd9 | |||
| db58cf2e68 | |||
| 1cf535865a | |||
| 6042d2a3c7 | |||
| df03819e12 | |||
| 072354afb7 | |||
| 57c7369aea | |||
| f92cdf335f | |||
| 19f3c857d5 | |||
| 7f3da73aa4 | |||
| c384c2b610 | |||
| 16d7dd426a | |||
| 787b6928d9 | |||
| ddbdbc5fb5 | |||
| d0edf9fb48 | |||
| 28d4446aa7 | |||
| ab3419ab09 | |||
| 457fe1535f | |||
| 6e90b02052 | |||
| fdc2680ae8 | |||
| 35a398443a | |||
| b55ce8aa93 | |||
| 74e67ef702 | |||
| e689248919 | |||
| 5a527b3fc9 | |||
| 565cc37912 | |||
| 50e1864a85 | |||
| 37ac88fc1b | |||
| 8ed98bcb07 | |||
| b8609ec149 | |||
| 96eb603bca | |||
| 20aaf35fc4 | |||
| 151a63a118 | |||
| 9db2267445 | |||
| 69a7ef063d | |||
| 8527875f0a | |||
| 2244da07e6 | |||
| 18bf710549 | |||
| 9e2e5a3f7f | |||
| 3df7fd6dd1 | |||
| 4371c46c4c | |||
| fcf7c1275b | |||
| e6e62fa767 | |||
| 8e1cb8fb70 | |||
| 00674686cb | |||
| b350bb2e7a | |||
| e17c704a4b | |||
| 3259245780 | |||
| 63f21084b6 | |||
| ced3555248 | |||
| 27eef24611 | |||
| dcf358f27d | |||
| 40afc8c0e2 | |||
| 1dabdbf7d8 | |||
| d65397a0ba | |||
| 8c72500070 | |||
| 2a27e21e00 | |||
| 3bc52a1f2c | |||
| 9806b5a0a9 | |||
| 201cd0f5fc | |||
| 5e2e92c11a | |||
| d57657599b | |||
| 995db834b2 | |||
| b22bc8a799 | |||
| dfd22aac46 | |||
| 72f62718f1 | |||
| e938a4d154 | |||
| 641019262e | |||
| fe9056f081 | |||
| fba49b8088 | |||
| a5ae1eea30 | |||
| 6e0ee9ec32 | |||
| a5e3e10e84 | |||
| 99afbdd606 | |||
| 90c0b83c34 | |||
| 68f9dd1ce1 | |||
| 51e4b1de99 | |||
| 08f091a889 | |||
| 692e3cc56e | |||
| bcfbccc381 | |||
| 8aa8c65d0e | |||
| 635d485b69 | |||
| c0630eeebb | |||
| 8932be2d49 | |||
| 3905486240 | |||
| eedc14b13c | |||
| 44bb615eee | |||
| 541d387143 | |||
| a8ad9c10bd | |||
| f1669224da | |||
| 2a27061891 | |||
| 6c55b2de58 | |||
| 9b99855c43 | |||
| 0269e88575 | |||
| 7fbd43283d | |||
| 226c123b77 | |||
| ac271d1006 | |||
| af84425689 | |||
| 512e9c947c | |||
| e7319376a5 | |||
| 2a7b493769 | |||
| f516a0d2e4 | |||
| 62f872122c | |||
| 89737d6e00 | |||
| 6a81d54e53 | |||
| c0062746eb | |||
| 809a904bc8 | |||
| 602d27c7b0 | |||
| aad61e876f | |||
| eb0e9947f2 | |||
| 23a09cff1b | |||
| ebe9041fdc | |||
| f93ef09b58 | |||
| e74cfb93b5 | |||
| 4a44621f87 | |||
| c7acaa2f5e | |||
| 139abad1f4 | |||
| a3a5306f11 | |||
| fb1c3bc446 | |||
| aaf344a4dd | |||
| 62ca9c0ed2 | |||
| dc8be8740d | |||
| d9bcf4df92 | |||
| 7ceb94f9c2 | |||
| 2e5becb4fb | |||
| 5e12f568bd | |||
| 80382c0bf9 | |||
| 91150d4150 | |||
| 6bfc38db53 | |||
| 95b99db199 | |||
| 1b13395e65 | |||
| fe21904b53 | |||
| ab0d666f03 | |||
| 30add7a765 | |||
| 83971a1913 | |||
| 2f62081683 | |||
| c7eb81dfa4 | |||
| 9f35f526e0 | |||
| e755a63250 | |||
| 29c6b62ba1 | |||
| 9d69903c36 | |||
| 51475a9290 | |||
| a9e794bde9 | |||
| 5114a7aa27 | |||
| d14042e536 | |||
| 7819fca349 | |||
| 68d9cfb550 | |||
| 1b7fd95214 | |||
| 0a1e6ccf9a | |||
| 0db3f415a8 | |||
| 8a1385b9d0 | |||
| a52143b0ef | |||
| 75ec41c85a | |||
| 827c8b3c48 | |||
| 194b35d889 | |||
| 1b33523537 | |||
| 807b95597a | |||
| 14b1ffa413 | |||
| d1db4d5534 | |||
| a45c0e537f | |||
| 4dab9b8fa3 | |||
| a84f8ba5d6 | |||
| f6f5cab661 | |||
| 618f563ce9 | |||
| 5b49c90538 | |||
| 41be0003f1 | |||
| 8f8ee28ba0 | |||
| b3ae7fbb49 | |||
| 837854de1e | |||
| 8cc1f0726f | |||
| e1d617ef70 | |||
| 5f199d68f9 | |||
| b8cca2db97 | |||
| 35e959219d | |||
| 08d466faee | |||
| 0b5823f451 | |||
| c77b150c28 | |||
| 3cf27bb838 | |||
| 26a90435c7 | |||
| f6efaba906 | |||
| ff0e831da9 | |||
| 96bb65723a | |||
| 33ac4bc424 | |||
| 698503b467 | |||
| 0657525d40 | |||
| 064d0de531 | |||
| 471bf36a7b | |||
| cb7d2b4040 | |||
| 6032cbcf45 | |||
| 73785d7552 | |||
| 431b5ffa59 | |||
| c0500a0d4d | |||
| 5300534188 | |||
| 02192a5f53 | |||
| 2a98d5b8ee | |||
| 4f495b5fc6 | |||
| b75e2d23a2 | |||
| 5261cdc794 | |||
| b179f61c6f | |||
| 71b245ad6f | |||
| 5b070cf87a | |||
| b8afe0b364 | |||
| 92b4ec48f7 | |||
| 6a69ac997d | |||
| 8c542c30a9 | |||
| a8388c841f |
@@ -0,0 +1,5 @@
|
||||
---
|
||||
"llamaindex": patch
|
||||
---
|
||||
|
||||
Fixed errors (#225 and #226) Thanks @marcusschiesser
|
||||
@@ -1,5 +0,0 @@
|
||||
---
|
||||
"llamaindex": patch
|
||||
---
|
||||
|
||||
Chat History summarization (thanks @marcusschlesser)
|
||||
@@ -1,5 +0,0 @@
|
||||
---
|
||||
"llamaindex": patch
|
||||
---
|
||||
|
||||
Notion database support (thanks @TomPenguin)
|
||||
@@ -1,5 +0,0 @@
|
||||
---
|
||||
"llamaindex": patch
|
||||
---
|
||||
|
||||
KeywordIndex (thanks @swk777)
|
||||
@@ -1,7 +1,6 @@
|
||||
name: Bugfix
|
||||
title: "Sweep: "
|
||||
title: ""
|
||||
description: Write something like "We notice ... behavior when ... happens instead of ...""
|
||||
labels: sweep
|
||||
body:
|
||||
- type: textarea
|
||||
id: description
|
||||
@@ -1,11 +1,10 @@
|
||||
name: Feature Request
|
||||
title: "Sweep: "
|
||||
description: Write something like "Write an api endpoint that does "..." in the "..." file"
|
||||
labels: sweep
|
||||
title: ""
|
||||
description: Write something like "Write an api endpoint that does "..." in the "..." file". If you would like to use sweep.dev prefix with "Sweep:"
|
||||
body:
|
||||
- type: textarea
|
||||
id: description
|
||||
attributes:
|
||||
label: Details
|
||||
description: More details for Sweep
|
||||
description: More details
|
||||
placeholder: The new endpoint should use the ... class from ... file because it contains ... logic
|
||||
@@ -1,11 +1,10 @@
|
||||
name: Refactor
|
||||
title: "Sweep: "
|
||||
description: Write something like "Modify the ... api endpoint to use ... version and ... framework"
|
||||
labels: sweep
|
||||
title: ""
|
||||
description: Write something like "Modify the ... api endpoint to use ... version and ... framework" If you would like to use sweep.dev prefix with "Sweep:"
|
||||
body:
|
||||
- type: textarea
|
||||
id: description
|
||||
attributes:
|
||||
label: Details
|
||||
description: More details for Sweep
|
||||
description: More details
|
||||
placeholder: We are migrating this function to ... version because ...
|
||||
@@ -3,6 +3,7 @@
|
||||
# dependencies
|
||||
node_modules
|
||||
.pnp
|
||||
.pnpm-store
|
||||
.pnp.js
|
||||
|
||||
# testing
|
||||
@@ -36,3 +37,6 @@ yarn-error.log*
|
||||
.vercel
|
||||
|
||||
dist/
|
||||
|
||||
# vs code
|
||||
.vscode/launch.json
|
||||
|
||||
@@ -2,3 +2,4 @@
|
||||
. "$(dirname -- "$0")/_/husky.sh"
|
||||
|
||||
pnpm lint
|
||||
npx lint-staged
|
||||
|
||||
Vendored
+3
-2
@@ -4,5 +4,6 @@
|
||||
"editor.defaultFormatter": "esbenp.prettier-vscode",
|
||||
"[xml]": {
|
||||
"editor.defaultFormatter": "redhat.vscode-xml"
|
||||
}
|
||||
}
|
||||
},
|
||||
"jest.rootPath": "./packages/core"
|
||||
}
|
||||
@@ -84,6 +84,26 @@ Check out our NextJS playground at https://llama-playground.vercel.app/. The sou
|
||||
|
||||
- [SimplePrompt](/packages/core/src/Prompt.ts): A simple standardized function call definition that takes in inputs and formats them in a template literal. SimplePrompts can be specialized using currying and combined using other SimplePrompt functions.
|
||||
|
||||
## Note: NextJS:
|
||||
|
||||
If you're using NextJS App Router, you'll need to use the NodeJS runtime (default) and add the follow config to your next.config.js to have it use imports/exports in the same way Node does.
|
||||
|
||||
```js
|
||||
export const runtime = "nodejs"; // default
|
||||
```
|
||||
|
||||
```js
|
||||
// next.config.js
|
||||
/** @type {import('next').NextConfig} */
|
||||
const nextConfig = {
|
||||
experimental: {
|
||||
serverComponentsExternalPackages: ["pdf-parse"], // Puts pdf-parse in actual NodeJS mode with NextJS App Router
|
||||
},
|
||||
};
|
||||
|
||||
module.exports = nextConfig;
|
||||
```
|
||||
|
||||
## Supported LLMs:
|
||||
|
||||
- OpenAI GPT-3.5-turbo and GPT-4
|
||||
|
||||
@@ -6,6 +6,8 @@ sidebar_position: 4
|
||||
|
||||
We include several end-to-end examples using LlamaIndex.TS in the repository
|
||||
|
||||
Check out the examples below or try them out and complete them in minutes with interactive Github Codespace tutorials provided by Dev-Docs [here](https://codespaces.new/team-dev-docs/lits-dev-docs-playground?devcontainer_path=.devcontainer%2Fjavascript_ltsquickstart%2Fdevcontainer.json):
|
||||
|
||||
## [Chat Engine](https://github.com/run-llama/LlamaIndexTS/blob/main/apps/simple/chatEngine.ts)
|
||||
|
||||
Read a file and chat about it with the LLM.
|
||||
@@ -14,7 +16,7 @@ Read a file and chat about it with the LLM.
|
||||
|
||||
Create a vector index and query it. The vector index will use embeddings to fetch the top k most relevant nodes. By default, the top k is 2.
|
||||
|
||||
## [Summary Index](https://github.com/run-llama/LlamaIndexTS/blob/main/apps/simple/summarIndex.ts)
|
||||
## [Summary Index](https://github.com/run-llama/LlamaIndexTS/blob/main/apps/simple/summaryIndex.ts)
|
||||
|
||||
Create a list index and query it. This example also use the `LLMRetriever`, which will use the LLM to select the best nodes to use when generating answer.
|
||||
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
---
|
||||
sidebar_position: 5
|
||||
---
|
||||
|
||||
# Environments
|
||||
|
||||
LlamaIndex currently officially supports NodeJS 18 and NodeJS 20.
|
||||
|
||||
## NextJS App Router
|
||||
|
||||
If you're using NextJS App Router route handlers/serverless functions, you'll need to use the NodeJS mode:
|
||||
|
||||
```js
|
||||
export const runtime = "nodejs"; // default
|
||||
```
|
||||
|
||||
and you'll need to add an exception for pdf-parse in your next.config.js
|
||||
|
||||
```js
|
||||
// next.config.js
|
||||
/** @type {import('next').NextConfig} */
|
||||
const nextConfig = {
|
||||
experimental: {
|
||||
serverComponentsExternalPackages: ["pdf-parse"], // Puts pdf-parse in actual NodeJS mode with NextJS App Router
|
||||
},
|
||||
};
|
||||
|
||||
module.exports = nextConfig;
|
||||
```
|
||||
@@ -19,7 +19,7 @@ That's where **LlamaIndex.TS** comes in.
|
||||
|
||||
LlamaIndex.TS provides the following tools:
|
||||
|
||||
- **Data loading** ingest your existing `txt` and `pdf` data directly
|
||||
- **Data loading** ingest your existing `.txt`, `.pdf`, `.csv`, `.md` and `.docx` data directly
|
||||
- **Data indexes** structure your data in intermediate representations that are easy and performant for LLMs to consume.
|
||||
- **Engines** provide natural language access to your data. For example:
|
||||
- Query engines are powerful retrieval interfaces for knowledge-augmented output.
|
||||
|
||||
@@ -4,7 +4,7 @@ sidebar_position: 1
|
||||
|
||||
# Reader / Loader
|
||||
|
||||
LlamaIndex.TS supports easy loading of files from folders using the `SimpleDirectoryReader` class. Currently, `.txt` and `.pdf` files are supported, with more planned in the future!
|
||||
LlamaIndex.TS supports easy loading of files from folders using the `SimpleDirectoryReader` class. Currently, `.txt`, `.pdf`, `.csv`, `.md` and `.docx` files are supported, with more planned in the future!
|
||||
|
||||
```typescript
|
||||
import { SimpleDirectoryReader } from "llamaindex";
|
||||
|
||||
@@ -15,24 +15,24 @@
|
||||
"typecheck": "tsc"
|
||||
},
|
||||
"dependencies": {
|
||||
"@docusaurus/core": "2.4.1",
|
||||
"@docusaurus/preset-classic": "2.4.1",
|
||||
"@docusaurus/remark-plugin-npm2yarn": "^2.4.1",
|
||||
"@docusaurus/core": "2.4.3",
|
||||
"@docusaurus/preset-classic": "2.4.3",
|
||||
"@docusaurus/remark-plugin-npm2yarn": "^2.4.3",
|
||||
"@mdx-js/react": "^1.6.22",
|
||||
"clsx": "^1.2.1",
|
||||
"postcss": "^8.4.28",
|
||||
"postcss": "^8.4.31",
|
||||
"prism-react-renderer": "^1.3.5",
|
||||
"raw-loader": "^4.0.2",
|
||||
"react": "^17.0.2",
|
||||
"react-dom": "^17.0.2"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@docusaurus/module-type-aliases": "2.4.1",
|
||||
"@docusaurus/types": "^2.4.1",
|
||||
"@tsconfig/docusaurus": "^1.0.7",
|
||||
"@docusaurus/module-type-aliases": "2.4.3",
|
||||
"@docusaurus/types": "^2.4.3",
|
||||
"@tsconfig/docusaurus": "^2.0.1",
|
||||
"docusaurus-plugin-typedoc": "^0.19.2",
|
||||
"typedoc": "^0.24.8",
|
||||
"typedoc-plugin-markdown": "^3.15.4",
|
||||
"typedoc-plugin-markdown": "^3.16.0",
|
||||
"typescript": "^4.9.5"
|
||||
},
|
||||
"browserslist": {
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import * as dotenv from "dotenv";
|
||||
import * as fs from "fs";
|
||||
import { MongoClient } from "mongodb";
|
||||
|
||||
// Load environment variables from local .env file
|
||||
dotenv.config();
|
||||
|
||||
const jsonFile = "tinytweets.json";
|
||||
const mongoUri = process.env.MONGODB_URI!;
|
||||
const databaseName = process.env.MONGODB_DATABASE!;
|
||||
const collectionName = process.env.MONGODB_COLLECTION!;
|
||||
|
||||
async function importJsonToMongo() {
|
||||
// Load the tweets from a local file
|
||||
const tweets = JSON.parse(fs.readFileSync(jsonFile, "utf-8"));
|
||||
|
||||
// Create a new client and connect to the server
|
||||
const client = new MongoClient(mongoUri);
|
||||
|
||||
const db = client.db(databaseName);
|
||||
const collection = db.collection(collectionName);
|
||||
|
||||
// Insert the tweets into mongo
|
||||
await collection.insertMany(tweets);
|
||||
|
||||
console.log(
|
||||
`Data imported successfully to the MongoDB collection ${collectionName}.`,
|
||||
);
|
||||
await client.close();
|
||||
}
|
||||
|
||||
// Run the import function
|
||||
importJsonToMongo();
|
||||
@@ -0,0 +1,50 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import * as dotenv from "dotenv";
|
||||
import {
|
||||
MongoDBAtlasVectorSearch,
|
||||
SimpleMongoReader,
|
||||
storageContextFromDefaults,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
import { MongoClient } from "mongodb";
|
||||
|
||||
// Load environment variables from local .env file
|
||||
dotenv.config();
|
||||
|
||||
const mongoUri = process.env.MONGODB_URI!;
|
||||
const databaseName = process.env.MONGODB_DATABASE!;
|
||||
const collectionName = process.env.MONGODB_COLLECTION!;
|
||||
const vectorCollectionName = process.env.MONGODB_VECTORS!;
|
||||
const indexName = process.env.MONGODB_VECTOR_INDEX!;
|
||||
|
||||
async function loadAndIndex() {
|
||||
// Create a new client and connect to the server
|
||||
const client = new MongoClient(mongoUri);
|
||||
// load objects from mongo and convert them into LlamaIndex Document objects
|
||||
// llamaindex has a special class that does this for you
|
||||
// it pulls every object in a given collection
|
||||
const reader = new SimpleMongoReader(client);
|
||||
const documents = await reader.loadData(databaseName, collectionName, [
|
||||
"full_text",
|
||||
]);
|
||||
|
||||
// create Atlas as a vector store
|
||||
const vectorStore = new MongoDBAtlasVectorSearch({
|
||||
mongodbClient: client,
|
||||
dbName: databaseName,
|
||||
collectionName: vectorCollectionName, // this is where your embeddings will be stored
|
||||
indexName: indexName, // this is the name of the index you will need to create
|
||||
});
|
||||
|
||||
// now create an index from all the Documents and store them in Atlas
|
||||
const storageContext = await storageContextFromDefaults({ vectorStore });
|
||||
await VectorStoreIndex.fromDocuments(documents, { storageContext });
|
||||
console.log(
|
||||
`Successfully created embeddings in the MongoDB collection ${vectorCollectionName}.`,
|
||||
);
|
||||
await client.close();
|
||||
}
|
||||
|
||||
loadAndIndex();
|
||||
|
||||
// you can't query your index yet because you need to create a vector search index in mongodb's UI now
|
||||
@@ -0,0 +1,34 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import * as dotenv from "dotenv";
|
||||
import {
|
||||
MongoDBAtlasVectorSearch,
|
||||
serviceContextFromDefaults,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
import { MongoClient } from "mongodb";
|
||||
|
||||
// Load environment variables from local .env file
|
||||
dotenv.config();
|
||||
|
||||
async function query() {
|
||||
const client = new MongoClient(process.env.MONGODB_URI!);
|
||||
const serviceContext = serviceContextFromDefaults();
|
||||
const store = new MongoDBAtlasVectorSearch({
|
||||
mongodbClient: client,
|
||||
dbName: process.env.MONGODB_DATABASE!,
|
||||
collectionName: process.env.MONGODB_VECTORS!,
|
||||
indexName: process.env.MONGODB_VECTOR_INDEX!,
|
||||
});
|
||||
|
||||
const index = await VectorStoreIndex.fromVectorStore(store, serviceContext);
|
||||
|
||||
const retriever = index.asRetriever({ similarityTopK: 20 });
|
||||
const queryEngine = index.asQueryEngine({ retriever });
|
||||
const result = await queryEngine.query(
|
||||
"What does the author think of web frameworks?",
|
||||
);
|
||||
console.log(result.response);
|
||||
await client.close();
|
||||
}
|
||||
|
||||
query();
|
||||
@@ -0,0 +1,13 @@
|
||||
# mongodb-llamaindexts
|
||||
|
||||
## 0.0.2
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- llamaindex@0.0.36
|
||||
@@ -0,0 +1,127 @@
|
||||
# LlamaIndexTS retrieval augmented generation with MongoDB
|
||||
|
||||
### Prepare Environment
|
||||
|
||||
Make sure to run `pnpm install` and set your OpenAI environment variable before running these examples.
|
||||
|
||||
```
|
||||
pnpm install
|
||||
export OPENAI_API_KEY="sk-..."
|
||||
```
|
||||
|
||||
### Sign up for MongoDB Atlas
|
||||
|
||||
We'll be using MongoDB's hosted database service, [MongoDB Atlas](https://www.mongodb.com/cloud/atlas/register). You can sign up for free and get a small hosted cluster for free:
|
||||
|
||||

|
||||
|
||||
The signup process will walk you through the process of creating your cluster and ensuring it's configured for you to access. Once the cluster is created, choose "Connect" and then "Connect to your application". Choose Python, and you'll be presented with a connection string that looks like this:
|
||||
|
||||

|
||||
|
||||
### Set up environment variables
|
||||
|
||||
Copy the connection string (make sure you include your password) and put it into a file called `.env` in the root of this repo. It should look like this:
|
||||
|
||||
```
|
||||
MONGODB_URI=mongodb+srv://seldo:xxxxxxxxxxx@llamaindexdemocluster.xfrdhpz.mongodb.net/?retryWrites=true&w=majority
|
||||
```
|
||||
|
||||
You will also need to choose a name for your database, and the collection where we will store the tweets, and also include them in .env. They can be any string, but this is what we used:
|
||||
|
||||
```
|
||||
MONGODB_DATABASE=tiny_tweets_db
|
||||
MONGODB_COLLECTION=tiny_tweets_collection
|
||||
```
|
||||
|
||||
### Import tweets into MongoDB
|
||||
|
||||
You are now ready to import our ready-made data set into Mongo. This is the file `tinytweets.json`, a selection of approximately 1000 tweets from @seldo on Twitter in mid-2019. With your environment set up you can do this by running
|
||||
|
||||
```
|
||||
pnpm ts-node 1_import.ts
|
||||
```
|
||||
|
||||
If you don't want to use tweets, you can replace `json_file` with any other array of JSON objects, but you will need to modify some code later to make sure the correct field gets indexed. There is no LlamaIndex-specific code here; you can load your data into Mongo any way you want to.
|
||||
|
||||
### Load and index your data
|
||||
|
||||
Now we're ready to index our data. To do this, LlamaIndex will pull your text out of Mongo, split it into chunks, and then send those chunks to OpenAI to be turned into [vector embeddings](https://docs.llamaindex.ai/en/stable/understanding/indexing/indexing.html#what-is-an-embedding). The embeddings will then be stored in a new collection in Mongo. This will take a while depending how much text you have, but the good news is that once it's done you will be able to query quickly without needing to re-index.
|
||||
|
||||
We'll be using OpenAI to do the embedding, so now is when you need to [generate an OpenAI API key](https://platform.openai.com/account/api-keys) if you haven't already and add it to your `.env` file like this:
|
||||
|
||||
```
|
||||
OPENAI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
|
||||
```
|
||||
|
||||
You'll also need to pick a name for the new collection where the embeddings will be stored, and add it to `.env`, along with the name of a vector search index (we'll be creating this in the next step, after you've indexed your data):
|
||||
|
||||
```
|
||||
MONGODB_VECTORS=tiny_tweets_vectors
|
||||
MONGODB_VECTOR_INDEX=tiny_tweets_vector_index
|
||||
```
|
||||
|
||||
If the data you're indexing is the tweets we gave you, you're ready to go:
|
||||
|
||||
```bash
|
||||
pnpm ts-node 2_load_and_index.ts
|
||||
```
|
||||
|
||||
> Note: this script is running a couple of minutes and currently doesn't show any progress.
|
||||
|
||||
What you're doing here is creating a Reader which loads the data out of Mongo in the collection and database specified. It looks for text in a set of specific keys in each object. In this case we've given it just one key, "full_text".
|
||||
|
||||
Now you're creating a vector search client for Mongo. In addition to a MongoDB client object, you again tell it what database everything is in. This time you give it the name of the collection where you'll store the vector embeddings, and the name of the vector search index you'll create in the next step.
|
||||
|
||||
### Create a vector search index
|
||||
|
||||
Now if all has gone well you should be able to log in to the Mongo Atlas UI and see two collections in your database: the original data in `tiny_tweets_collection`, and the vector embeddings in `tiny_tweets_vectors`.
|
||||
|
||||

|
||||
|
||||
Now it's time to create the vector search index so that you can query the data.
|
||||
It's not yet possible to programmatically create a vector search index using the [`createIndex`](https://www.mongodb.com/docs/manual/reference/method/db.collection.createIndex/) function, therefore we have to create one manually in the UI.
|
||||
To do so, first, click the Search tab, and then click "Create Search Index":
|
||||
|
||||

|
||||
|
||||
We have to use the JSON editor, as the Visual Editor does not yet support to create a vector search index:
|
||||
|
||||

|
||||
|
||||
Now under "database and collection" select `tiny_tweets_db` and within that select `tiny_tweets_vectors`. Then under "Index name" enter `tiny_tweets_vector_index` (or whatever value you put for MONGODB_VECTOR_INDEX in `.env`). Under that, you'll want to enter this JSON object:
|
||||
|
||||
```json
|
||||
{
|
||||
"mappings": {
|
||||
"dynamic": true,
|
||||
"fields": {
|
||||
"embedding": {
|
||||
"dimensions": 1536,
|
||||
"similarity": "cosine",
|
||||
"type": "knnVector"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
This tells Mongo that the `embedding` field in each document (in the `tiny_tweets_vectors` collection) is a vector of 1536 dimensions (this is the size of embeddings used by OpenAI), and that we want to use cosine similarity to compare vectors. You don't need to worry too much about these values unless you want to use a different LLM to OpenAI entirely.
|
||||
|
||||
The UI will ask you to review and confirm your choices, then you need to wait a minute or two while it generates the index. If all goes well, you should see something like this screen:
|
||||
|
||||

|
||||
|
||||
Now you're ready to query your data!
|
||||
|
||||
### Run a test query
|
||||
|
||||
You can do this by running
|
||||
|
||||
```bash
|
||||
pnpm ts-node 3_query.ts
|
||||
```
|
||||
|
||||
This sets up a connection to Atlas just like `2_load_and_index.ts` did, then it creates a [query engine](https://docs.llamaindex.ai/en/stable/understanding/querying/querying.html#getting-started) and runs a query against it.
|
||||
|
||||
If all is well, you should get a nuanced opinion about web frameworks.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 141 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 107 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 360 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 230 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 278 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 211 KiB |
@@ -0,0 +1,17 @@
|
||||
{
|
||||
"version": "0.0.2",
|
||||
"private": true,
|
||||
"name": "mongodb-llamaindexts",
|
||||
"dependencies": {
|
||||
"llamaindex": "workspace:*",
|
||||
"dotenv": "^16.3.1",
|
||||
"mongodb": "^6.2.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^18.18.6",
|
||||
"ts-node": "^10.9.1"
|
||||
},
|
||||
"scripts": {
|
||||
"lint": "eslint ."
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,5 +1,93 @@
|
||||
# simple
|
||||
|
||||
## 0.0.34
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- Updated dependencies
|
||||
- llamaindex@0.0.36
|
||||
|
||||
## 0.0.33
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [63f2108]
|
||||
- llamaindex@0.0.35
|
||||
|
||||
## 0.0.32
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [2a27e21]
|
||||
- llamaindex@0.0.34
|
||||
|
||||
## 0.0.31
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [5e2e92c]
|
||||
- llamaindex@0.0.33
|
||||
|
||||
## 0.0.30
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [90c0b83]
|
||||
- Updated dependencies [dfd22aa]
|
||||
- llamaindex@0.0.32
|
||||
|
||||
## 0.0.29
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [6c55b2d]
|
||||
- Updated dependencies [8aa8c65]
|
||||
- Updated dependencies [6c55b2d]
|
||||
- llamaindex@0.0.31
|
||||
|
||||
## 0.0.28
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [139abad]
|
||||
- Updated dependencies [139abad]
|
||||
- Updated dependencies [eb0e994]
|
||||
- Updated dependencies [eb0e994]
|
||||
- Updated dependencies [139abad]
|
||||
- llamaindex@0.0.30
|
||||
|
||||
## 0.0.27
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [a52143b]
|
||||
- Updated dependencies [1b7fd95]
|
||||
- Updated dependencies [0db3f41]
|
||||
- llamaindex@0.0.29
|
||||
|
||||
## 0.0.26
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [96bb657]
|
||||
- Updated dependencies [96bb657]
|
||||
- Updated dependencies [837854d]
|
||||
- llamaindex@0.0.28
|
||||
|
||||
## 0.0.25
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Updated dependencies [4a5591b]
|
||||
- Updated dependencies [4a5591b]
|
||||
- Updated dependencies [4a5591b]
|
||||
- llamaindex@0.0.27
|
||||
|
||||
## 0.0.24
|
||||
|
||||
### Patch Changes
|
||||
|
||||
@@ -4,8 +4,6 @@ import { Anthropic } from "llamaindex";
|
||||
const anthropic = new Anthropic();
|
||||
const result = await anthropic.chat([
|
||||
{ content: "You want to talk in rhymes.", role: "system" },
|
||||
{ content: "Hello, world!", role: "user" },
|
||||
{ content: "Hello!", role: "assistant" },
|
||||
{
|
||||
content:
|
||||
"How much wood would a woodchuck chuck if a woodchuck could chuck wood?",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,24 @@
|
||||
import { SimpleDirectoryReader } from "llamaindex";
|
||||
|
||||
function callback(
|
||||
category: string,
|
||||
name: string,
|
||||
status: any,
|
||||
message?: string,
|
||||
): boolean {
|
||||
console.log(category, name, status, message);
|
||||
if (name.endsWith(".pdf")) {
|
||||
console.log("I DON'T WANT PDF FILES!");
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
async function main() {
|
||||
// Load page
|
||||
const reader = new SimpleDirectoryReader(callback);
|
||||
const params = { directoryPath: "./data" };
|
||||
await reader.loadData(params);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,21 @@
|
||||
import { HTMLReader, VectorStoreIndex } from "llamaindex";
|
||||
|
||||
async function main() {
|
||||
// Load page
|
||||
const reader = new HTMLReader();
|
||||
const documents = await reader.loadData("data/18-1_Changelog.html");
|
||||
|
||||
// Split text and create embeddings. Store them in a VectorStoreIndex
|
||||
const index = await VectorStoreIndex.fromDocuments(documents);
|
||||
|
||||
// Query the index
|
||||
const queryEngine = index.asQueryEngine();
|
||||
const response = await queryEngine.query(
|
||||
"What were the notable changes in 18.1?",
|
||||
);
|
||||
|
||||
// Output response
|
||||
console.log(response.toString());
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,68 @@
|
||||
import { MongoClient } from "mongodb";
|
||||
import { VectorStoreIndex } from "../../packages/core/src/indices";
|
||||
import { Document } from "../../packages/core/src/Node";
|
||||
import { SimpleMongoReader } from "../../packages/core/src/readers/SimpleMongoReader";
|
||||
|
||||
import { stdin as input, stdout as output } from "node:process";
|
||||
import readline from "node:readline/promises";
|
||||
|
||||
async function main() {
|
||||
//Dummy test code
|
||||
const query: object = { _id: "waldo" };
|
||||
const options: object = {};
|
||||
const projections: object = { embedding: 0 };
|
||||
const limit: number = Infinity;
|
||||
const uri: string = process.env.MONGODB_URI ?? "fake_uri";
|
||||
const client: MongoClient = new MongoClient(uri);
|
||||
|
||||
//Where the real code starts
|
||||
const MR = new SimpleMongoReader(client);
|
||||
const documents: Document[] = await MR.loadData(
|
||||
"data",
|
||||
"posts",
|
||||
1,
|
||||
{},
|
||||
options,
|
||||
projections,
|
||||
);
|
||||
|
||||
//
|
||||
//If you need to look at low-level details of
|
||||
// a queryEngine (for example, needing to check each individual node)
|
||||
//
|
||||
|
||||
// Split text and create embeddings. Store them in a VectorStoreIndex
|
||||
// var storageContext = await storageContextFromDefaults({});
|
||||
// var serviceContext = serviceContextFromDefaults({});
|
||||
// const docStore = storageContext.docStore;
|
||||
|
||||
// for (const doc of documents) {
|
||||
// docStore.setDocumentHash(doc.id_, doc.hash);
|
||||
// }
|
||||
// const nodes = serviceContext.nodeParser.getNodesFromDocuments(documents);
|
||||
// console.log(nodes);
|
||||
|
||||
//
|
||||
//Making Vector Store from documents
|
||||
//
|
||||
|
||||
const index = await VectorStoreIndex.fromDocuments(documents);
|
||||
// Create query engine
|
||||
const queryEngine = index.asQueryEngine();
|
||||
|
||||
const rl = readline.createInterface({ input, output });
|
||||
while (true) {
|
||||
const query = await rl.question("Query: ");
|
||||
|
||||
if (!query) {
|
||||
break;
|
||||
}
|
||||
|
||||
const response = await queryEngine.query(query);
|
||||
|
||||
// Output response
|
||||
console.log(response.toString());
|
||||
}
|
||||
}
|
||||
|
||||
main();
|
||||
@@ -1,7 +1,7 @@
|
||||
import { OpenAI } from "llamaindex";
|
||||
|
||||
(async () => {
|
||||
const llm = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.0 });
|
||||
const llm = new OpenAI({ model: "gpt-4-1106-preview", temperature: 0.1 });
|
||||
|
||||
// complete api
|
||||
const response1 = await llm.complete("How are you?");
|
||||
@@ -9,7 +9,7 @@ import { OpenAI } from "llamaindex";
|
||||
|
||||
// chat api
|
||||
const response2 = await llm.chat([
|
||||
{ content: "Tell me a joke!", role: "user" },
|
||||
{ content: "Tell me a joke.", role: "user" },
|
||||
]);
|
||||
console.log(response2.message.content);
|
||||
})();
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
{
|
||||
"version": "0.0.24",
|
||||
"version": "0.0.34",
|
||||
"private": true,
|
||||
"name": "simple",
|
||||
"dependencies": {
|
||||
"@notionhq/client": "^2.2.12",
|
||||
"commander": "^11.0.0",
|
||||
"@notionhq/client": "^2.2.13",
|
||||
"@pinecone-database/pinecone": "^1.1.2",
|
||||
"commander": "^11.1.0",
|
||||
"llamaindex": "workspace:*"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^18.17.12"
|
||||
"@types/node": "^18.18.6",
|
||||
"ts-node": "^10.9.1"
|
||||
},
|
||||
"scripts": {
|
||||
"lint": "eslint ."
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
# Postgres Vector Store
|
||||
|
||||
There are two scripts available here: load-docs.ts and query.ts
|
||||
|
||||
## Prerequisites
|
||||
|
||||
You'll need a postgres database instance against which to run these scripts. A simple docker command would look like this:
|
||||
|
||||
> `docker run -d --rm --name vector-db -p 5432:5432 -e "POSTGRES_HOST_AUTH_METHOD=trust" ankane/pgvector`
|
||||
|
||||
Set the PGHOST and PGUSER (and PGPASSWORD) environment variables to match your database setup.
|
||||
|
||||
You'll also need a value for OPENAI_API_KEY in your environment.
|
||||
|
||||
**NOTE:** Using `--rm` in the example docker command above means that the vector store will be deleted every time the container is stopped. For production purposes, use a volume to ensure persistence across restarts.
|
||||
|
||||
## Setup and Loading Docs
|
||||
|
||||
Read and follow the instructions in the README.md file located one directory up to make sure your JS/TS dependencies are set up. The commands listed below are also run from that parent directory.
|
||||
|
||||
To import documents and save the embedding vectors to your database:
|
||||
|
||||
> `npx ts-node pg-vector-store/load-docs.ts data`
|
||||
|
||||
where data is the directory containing your input files. Using the _data_ directory in the example above will read all of the files in that directory using the llamaindexTS default readers for each file type.
|
||||
|
||||
## RAG Querying
|
||||
|
||||
To query using the resulting vector store:
|
||||
|
||||
> `npx ts-node pg-vector-store/query.ts`
|
||||
|
||||
The script will prompt for a question, then process and present the answer using the PGVectorStore data and your OpenAI API key. It will continue to prompt until you enter `q`, `quit` or `exit` as the next query.
|
||||
Executable
+68
@@ -0,0 +1,68 @@
|
||||
// load-docs.ts
|
||||
import fs from "fs/promises";
|
||||
import {
|
||||
SimpleDirectoryReader,
|
||||
storageContextFromDefaults,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
import { PGVectorStore } from "../../../packages/core/src/storage/vectorStore/PGVectorStore";
|
||||
|
||||
async function getSourceFilenames(sourceDir: string) {
|
||||
return await fs
|
||||
.readdir(sourceDir)
|
||||
.then((fileNames) => fileNames.map((file) => sourceDir + "/" + file));
|
||||
}
|
||||
|
||||
function callback(
|
||||
category: string,
|
||||
name: string,
|
||||
status: any,
|
||||
message: string = "",
|
||||
): boolean {
|
||||
console.log(category, name, status, message);
|
||||
return true;
|
||||
}
|
||||
|
||||
async function main(args: any) {
|
||||
const sourceDir: string = args.length > 2 ? args[2] : "../data";
|
||||
|
||||
console.log(`Finding documents in ${sourceDir}`);
|
||||
const fileList = await getSourceFilenames(sourceDir);
|
||||
const count = fileList.length;
|
||||
console.log(`Found ${count} files`);
|
||||
|
||||
console.log(`Importing contents from ${count} files in ${sourceDir}`);
|
||||
var fileName = "";
|
||||
try {
|
||||
// Passing callback fn to the ctor here
|
||||
// will enable looging to console.
|
||||
// See callback fn, defined above.
|
||||
const rdr = new SimpleDirectoryReader(callback);
|
||||
const docs = await rdr.loadData({ directoryPath: sourceDir });
|
||||
|
||||
const pgvs = new PGVectorStore();
|
||||
pgvs.setCollection(sourceDir);
|
||||
pgvs.clearCollection();
|
||||
|
||||
const ctx = await storageContextFromDefaults({ vectorStore: pgvs });
|
||||
|
||||
console.debug(" - creating vector store");
|
||||
const index = await VectorStoreIndex.fromDocuments(docs, {
|
||||
storageContext: ctx,
|
||||
});
|
||||
console.debug(" - done.");
|
||||
} catch (err) {
|
||||
console.error(fileName, err);
|
||||
console.log(
|
||||
"If your PGVectorStore init failed, make sure to set env vars for PGUSER or USER, PGHOST, PGPORT and PGPASSWORD as needed.",
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
console.log(
|
||||
"Done. Try running query.ts to ask questions against the imported embeddings.",
|
||||
);
|
||||
process.exit(0);
|
||||
}
|
||||
|
||||
main(process.argv).catch((err) => console.error(err));
|
||||
Executable
+67
@@ -0,0 +1,67 @@
|
||||
import { VectorStoreIndex } from "../../../packages/core/src/indices/vectorStore/VectorStoreIndex";
|
||||
import { serviceContextFromDefaults } from "../../../packages/core/src/ServiceContext";
|
||||
import { PGVectorStore } from "../../../packages/core/src/storage/vectorStore/PGVectorStore";
|
||||
|
||||
async function main() {
|
||||
const readline = require("readline").createInterface({
|
||||
input: process.stdin,
|
||||
output: process.stdout,
|
||||
});
|
||||
|
||||
try {
|
||||
const pgvs = new PGVectorStore();
|
||||
// Optional - set your collection name, default is no filter on this field.
|
||||
// pgvs.setCollection();
|
||||
|
||||
const ctx = serviceContextFromDefaults();
|
||||
const index = await VectorStoreIndex.fromVectorStore(pgvs, ctx);
|
||||
|
||||
// Query the index
|
||||
const queryEngine = await index.asQueryEngine();
|
||||
|
||||
let question = "";
|
||||
while (!isQuit(question)) {
|
||||
question = await getUserInput(readline);
|
||||
|
||||
if (isQuit(question)) {
|
||||
readline.close();
|
||||
process.exit(0);
|
||||
}
|
||||
|
||||
try {
|
||||
const answer = await queryEngine.query(question);
|
||||
console.log(answer.response);
|
||||
} catch (error) {
|
||||
console.error("Error:", error);
|
||||
}
|
||||
}
|
||||
} catch (err) {
|
||||
console.error(err);
|
||||
console.log(
|
||||
"If your PGVectorStore init failed, make sure to set env vars for PGUSER or USER, PGHOST, PGPORT and PGPASSWORD as needed.",
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
function isQuit(question: string) {
|
||||
return ["q", "quit", "exit"].includes(question.trim().toLowerCase());
|
||||
}
|
||||
|
||||
// Function to get user input as a promise
|
||||
function getUserInput(readline: any): Promise<string> {
|
||||
return new Promise((resolve) => {
|
||||
readline.question(
|
||||
"What would you like to know?\n>",
|
||||
(userInput: string) => {
|
||||
resolve(userInput);
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
main()
|
||||
.catch(console.error)
|
||||
.finally(() => {
|
||||
process.exit(1);
|
||||
});
|
||||
@@ -0,0 +1,23 @@
|
||||
import { Portkey } from "llamaindex";
|
||||
|
||||
(async () => {
|
||||
const llms = [{}];
|
||||
const portkey = new Portkey({
|
||||
mode: "single",
|
||||
llms: [
|
||||
{
|
||||
provider: "anyscale",
|
||||
virtual_key: "anyscale-3b3c04",
|
||||
model: "meta-llama/Llama-2-13b-chat-hf",
|
||||
max_tokens: 2000,
|
||||
},
|
||||
],
|
||||
});
|
||||
const result = portkey.stream_chat([
|
||||
{ role: "system", content: "You are a helpful assistant." },
|
||||
{ role: "user", content: "Tell me a joke." },
|
||||
]);
|
||||
for await (const res of result) {
|
||||
process.stdout.write(res);
|
||||
}
|
||||
})();
|
||||
@@ -2,7 +2,10 @@ import fs from "node:fs/promises";
|
||||
|
||||
import {
|
||||
Anthropic,
|
||||
anthropicTextQaPrompt,
|
||||
CompactAndRefine,
|
||||
Document,
|
||||
ResponseSynthesizer,
|
||||
serviceContextFromDefaults,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
@@ -18,12 +21,20 @@ async function main() {
|
||||
|
||||
// Split text and create embeddings. Store them in a VectorStoreIndex
|
||||
const serviceContext = serviceContextFromDefaults({ llm: new Anthropic() });
|
||||
|
||||
const responseSynthesizer = new ResponseSynthesizer({
|
||||
responseBuilder: new CompactAndRefine(
|
||||
serviceContext,
|
||||
anthropicTextQaPrompt,
|
||||
),
|
||||
});
|
||||
|
||||
const index = await VectorStoreIndex.fromDocuments([document], {
|
||||
serviceContext,
|
||||
});
|
||||
|
||||
// Query the index
|
||||
const queryEngine = index.asQueryEngine();
|
||||
const queryEngine = index.asQueryEngine({ responseSynthesizer });
|
||||
const response = await queryEngine.query(
|
||||
"What did the author do in college?",
|
||||
);
|
||||
|
||||
@@ -3,6 +3,7 @@ import {
|
||||
OpenAI,
|
||||
RetrieverQueryEngine,
|
||||
serviceContextFromDefaults,
|
||||
SimilarityPostprocessor,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
import essay from "./essay";
|
||||
@@ -21,8 +22,16 @@ async function main() {
|
||||
|
||||
const retriever = index.asRetriever();
|
||||
retriever.similarityTopK = 5;
|
||||
const nodePostprocessor = new SimilarityPostprocessor({
|
||||
similarityCutoff: 0.7,
|
||||
});
|
||||
// TODO: cannot pass responseSynthesizer into retriever query engine
|
||||
const queryEngine = new RetrieverQueryEngine(retriever);
|
||||
const queryEngine = new RetrieverQueryEngine(
|
||||
retriever,
|
||||
undefined,
|
||||
undefined,
|
||||
[nodePostprocessor],
|
||||
);
|
||||
|
||||
const response = await queryEngine.query(
|
||||
"What did the author do growing up?",
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
import {
|
||||
OpenAI,
|
||||
ResponseSynthesizer,
|
||||
RetrieverQueryEngine,
|
||||
serviceContextFromDefaults,
|
||||
TextNode,
|
||||
TreeSummarize,
|
||||
VectorIndexRetriever,
|
||||
VectorStore,
|
||||
VectorStoreIndex,
|
||||
VectorStoreQuery,
|
||||
VectorStoreQueryResult,
|
||||
} from "llamaindex";
|
||||
|
||||
import { Index, Pinecone, RecordMetadata } from "@pinecone-database/pinecone";
|
||||
|
||||
/**
|
||||
* Please do not use this class in production; it's only for demonstration purposes.
|
||||
*/
|
||||
class PineconeVectorStore<T extends RecordMetadata = RecordMetadata>
|
||||
implements VectorStore
|
||||
{
|
||||
storesText = true;
|
||||
isEmbeddingQuery = false;
|
||||
|
||||
indexName!: string;
|
||||
pineconeClient!: Pinecone;
|
||||
index!: Index<T>;
|
||||
|
||||
constructor({ indexName, client }: { indexName: string; client: Pinecone }) {
|
||||
this.indexName = indexName;
|
||||
this.pineconeClient = client;
|
||||
this.index = client.index<T>(indexName);
|
||||
}
|
||||
|
||||
client() {
|
||||
return this.pineconeClient;
|
||||
}
|
||||
|
||||
async query(
|
||||
query: VectorStoreQuery,
|
||||
kwargs?: any,
|
||||
): Promise<VectorStoreQueryResult> {
|
||||
let queryEmbedding: number[] = [];
|
||||
if (query.queryEmbedding) {
|
||||
if (typeof query.alpha === "number") {
|
||||
const alpha = query.alpha;
|
||||
queryEmbedding = query.queryEmbedding.map((v) => v * alpha);
|
||||
} else {
|
||||
queryEmbedding = query.queryEmbedding;
|
||||
}
|
||||
}
|
||||
|
||||
// Current LlamaIndexTS implementation only support exact match filter, so we use kwargs instead.
|
||||
const filter = kwargs?.filter || {};
|
||||
|
||||
const response = await this.index.query({
|
||||
filter,
|
||||
vector: queryEmbedding,
|
||||
topK: query.similarityTopK,
|
||||
includeValues: true,
|
||||
includeMetadata: true,
|
||||
});
|
||||
|
||||
console.log(
|
||||
`Numbers of vectors returned by Pinecone after preFilters are applied: ${
|
||||
response?.matches?.length || 0
|
||||
}.`,
|
||||
);
|
||||
|
||||
const topKIds: string[] = [];
|
||||
const topKNodes: TextNode[] = [];
|
||||
const topKScores: number[] = [];
|
||||
|
||||
const metadataToNode = (metadata?: T): Partial<TextNode> => {
|
||||
if (!metadata) {
|
||||
throw new Error("metadata is undefined.");
|
||||
}
|
||||
|
||||
const nodeContent = metadata["_node_content"];
|
||||
if (!nodeContent) {
|
||||
throw new Error("nodeContent is undefined.");
|
||||
}
|
||||
|
||||
if (typeof nodeContent !== "string") {
|
||||
throw new Error("nodeContent is not a string.");
|
||||
}
|
||||
|
||||
return JSON.parse(nodeContent);
|
||||
};
|
||||
|
||||
if (response.matches) {
|
||||
for (const match of response.matches) {
|
||||
const node = new TextNode({
|
||||
...metadataToNode(match.metadata),
|
||||
embedding: match.values,
|
||||
});
|
||||
|
||||
topKIds.push(match.id);
|
||||
topKNodes.push(node);
|
||||
topKScores.push(match.score ?? 0);
|
||||
}
|
||||
}
|
||||
|
||||
const result = {
|
||||
ids: topKIds,
|
||||
nodes: topKNodes,
|
||||
similarities: topKScores,
|
||||
};
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
add(): Promise<string[]> {
|
||||
return Promise.resolve([]);
|
||||
}
|
||||
|
||||
delete(): Promise<void> {
|
||||
throw new Error("Method `delete` not implemented.");
|
||||
}
|
||||
|
||||
persist(): Promise<void> {
|
||||
throw new Error("Method `persist` not implemented.");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The goal of this example is to show how to use Pinecone as a vector store
|
||||
* for LlamaIndexTS with(out) preFilters.
|
||||
*
|
||||
* It should not be used in production like that,
|
||||
* as you might want to find a proper PineconeVectorStore implementation.
|
||||
*/
|
||||
async function main() {
|
||||
process.env.PINECONE_API_KEY = "Your Pinecone API Key.";
|
||||
process.env.PINECONE_ENVIRONMENT = "Your Pinecone Environment.";
|
||||
process.env.PINECONE_PROJECT_ID = "Your Pinecone Project ID.";
|
||||
process.env.PINECONE_INDEX_NAME = "Your Pinecone Index Name.";
|
||||
process.env.OPENAI_API_KEY = "Your OpenAI API Key.";
|
||||
process.env.OPENAI_API_ORGANIZATION = "Your OpenAI API Organization.";
|
||||
|
||||
const getPineconeVectorStore = async () => {
|
||||
return new PineconeVectorStore({
|
||||
indexName: process.env.PINECONE_INDEX_NAME || "index-name",
|
||||
client: new Pinecone(),
|
||||
});
|
||||
};
|
||||
|
||||
const getServiceContext = () => {
|
||||
const openAI = new OpenAI({
|
||||
model: "gpt-4",
|
||||
apiKey: process.env.OPENAI_API_KEY,
|
||||
});
|
||||
|
||||
return serviceContextFromDefaults({
|
||||
llm: openAI,
|
||||
});
|
||||
};
|
||||
|
||||
const getQueryEngine = async (filter: unknown) => {
|
||||
const vectorStore = await getPineconeVectorStore();
|
||||
const serviceContext = getServiceContext();
|
||||
|
||||
const vectorStoreIndex = await VectorStoreIndex.fromVectorStore(
|
||||
vectorStore,
|
||||
serviceContext,
|
||||
);
|
||||
|
||||
const retriever = new VectorIndexRetriever({
|
||||
index: vectorStoreIndex,
|
||||
similarityTopK: 500,
|
||||
});
|
||||
|
||||
const responseSynthesizer = new ResponseSynthesizer({
|
||||
serviceContext,
|
||||
responseBuilder: new TreeSummarize(serviceContext),
|
||||
});
|
||||
|
||||
return new RetrieverQueryEngine(retriever, responseSynthesizer, {
|
||||
filter,
|
||||
});
|
||||
};
|
||||
|
||||
// whatever is a key from your metadata
|
||||
const queryEngine = await getQueryEngine({
|
||||
whatever: {
|
||||
$gte: 1,
|
||||
$lte: 100,
|
||||
},
|
||||
});
|
||||
|
||||
const response = await queryEngine.query("How many results do you have?");
|
||||
|
||||
console.log(response.toString());
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,15 @@
|
||||
import { OpenAI } from "llamaindex";
|
||||
|
||||
(async () => {
|
||||
const llm = new OpenAI({ model: "gpt-4-vision-preview", temperature: 0.1 });
|
||||
|
||||
// complete api
|
||||
const response1 = await llm.complete("How are you?");
|
||||
console.log(response1.message.content);
|
||||
|
||||
// chat api
|
||||
const response2 = await llm.chat([
|
||||
{ content: "Tell me a joke!", role: "user" },
|
||||
]);
|
||||
console.log(response2.message.content);
|
||||
})();
|
||||
@@ -4,8 +4,6 @@ import { Anthropic } from "llamaindex";
|
||||
const anthropic = new Anthropic();
|
||||
const result = await anthropic.chat([
|
||||
{ content: "You want to talk in rhymes.", role: "system" },
|
||||
{ content: "Hello, world!", role: "user" },
|
||||
{ content: "Hello!", role: "assistant" },
|
||||
{
|
||||
content:
|
||||
"How much wood would a woodchuck chuck if a woodchuck could chuck wood?",
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
import { ClipEmbedding, similarity, SimilarityType } from "llamaindex";
|
||||
|
||||
async function main() {
|
||||
const clip = new ClipEmbedding();
|
||||
|
||||
// Get text embeddings
|
||||
const text1 = "a car";
|
||||
const textEmbedding1 = await clip.getTextEmbedding(text1);
|
||||
const text2 = "a football match";
|
||||
const textEmbedding2 = await clip.getTextEmbedding(text2);
|
||||
|
||||
// Get image embedding
|
||||
const image =
|
||||
"https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/football-match.jpg";
|
||||
const imageEmbedding = await clip.getImageEmbedding(image);
|
||||
|
||||
// Calc similarity
|
||||
const sim1 = similarity(
|
||||
textEmbedding1,
|
||||
imageEmbedding,
|
||||
SimilarityType.DEFAULT,
|
||||
);
|
||||
const sim2 = similarity(
|
||||
textEmbedding2,
|
||||
imageEmbedding,
|
||||
SimilarityType.DEFAULT,
|
||||
);
|
||||
|
||||
console.log(`Similarity between "${text1}" and the image is ${sim1}`);
|
||||
console.log(`Similarity between "${text2}" and the image is ${sim2}`);
|
||||
}
|
||||
|
||||
main();
|
||||
@@ -0,0 +1,24 @@
|
||||
import { SimpleDirectoryReader } from "llamaindex";
|
||||
|
||||
function callback(
|
||||
category: string,
|
||||
name: string,
|
||||
status: any,
|
||||
message?: string,
|
||||
): boolean {
|
||||
console.log(category, name, status, message);
|
||||
if (name.endsWith(".pdf")) {
|
||||
console.log("I DON'T WANT PDF FILES!");
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
async function main() {
|
||||
// Load page
|
||||
const reader = new SimpleDirectoryReader(callback);
|
||||
const params = { directoryPath: "./data" };
|
||||
await reader.loadData(params);
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,21 @@
|
||||
import { HTMLReader, VectorStoreIndex } from "llamaindex";
|
||||
|
||||
async function main() {
|
||||
// Load page
|
||||
const reader = new HTMLReader();
|
||||
const documents = await reader.loadData("data/18-1_Changelog.html");
|
||||
|
||||
// Split text and create embeddings. Store them in a VectorStoreIndex
|
||||
const index = await VectorStoreIndex.fromDocuments(documents);
|
||||
|
||||
// Query the index
|
||||
const queryEngine = index.asQueryEngine();
|
||||
const response = await queryEngine.query(
|
||||
"What were the notable changes in 18.1?",
|
||||
);
|
||||
|
||||
// Output response
|
||||
console.log(response.toString());
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,68 @@
|
||||
import { MongoClient } from "mongodb";
|
||||
import { Document } from "../../packages/core/src/Node";
|
||||
import { VectorStoreIndex } from "../../packages/core/src/indices";
|
||||
import { SimpleMongoReader } from "../../packages/core/src/readers/SimpleMongoReader";
|
||||
|
||||
import { stdin as input, stdout as output } from "node:process";
|
||||
import readline from "node:readline/promises";
|
||||
|
||||
async function main() {
|
||||
//Dummy test code
|
||||
const query: object = { _id: "waldo" };
|
||||
const options: object = {};
|
||||
const projections: object = { embedding: 0 };
|
||||
const limit: number = Infinity;
|
||||
const uri: string = process.env.MONGODB_URI ?? "fake_uri";
|
||||
const client: MongoClient = new MongoClient(uri);
|
||||
|
||||
//Where the real code starts
|
||||
const MR = new SimpleMongoReader(client);
|
||||
const documents: Document[] = await MR.loadData(
|
||||
"data",
|
||||
"posts",
|
||||
1,
|
||||
{},
|
||||
options,
|
||||
projections,
|
||||
);
|
||||
|
||||
//
|
||||
//If you need to look at low-level details of
|
||||
// a queryEngine (for example, needing to check each individual node)
|
||||
//
|
||||
|
||||
// Split text and create embeddings. Store them in a VectorStoreIndex
|
||||
// var storageContext = await storageContextFromDefaults({});
|
||||
// var serviceContext = serviceContextFromDefaults({});
|
||||
// const docStore = storageContext.docStore;
|
||||
|
||||
// for (const doc of documents) {
|
||||
// docStore.setDocumentHash(doc.id_, doc.hash);
|
||||
// }
|
||||
// const nodes = serviceContext.nodeParser.getNodesFromDocuments(documents);
|
||||
// console.log(nodes);
|
||||
|
||||
//
|
||||
//Making Vector Store from documents
|
||||
//
|
||||
|
||||
const index = await VectorStoreIndex.fromDocuments(documents);
|
||||
// Create query engine
|
||||
const queryEngine = index.asQueryEngine();
|
||||
|
||||
const rl = readline.createInterface({ input, output });
|
||||
while (true) {
|
||||
const query = await rl.question("Query: ");
|
||||
|
||||
if (!query) {
|
||||
break;
|
||||
}
|
||||
|
||||
const response = await queryEngine.query(query);
|
||||
|
||||
// Output response
|
||||
console.log(response.toString());
|
||||
}
|
||||
}
|
||||
|
||||
main();
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
import { OpenAI } from "llamaindex";
|
||||
|
||||
(async () => {
|
||||
const llm = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.0 });
|
||||
const llm = new OpenAI({ model: "gpt-4-1106-preview", temperature: 0.1 });
|
||||
|
||||
// complete api
|
||||
const response1 = await llm.complete("How are you?");
|
||||
@@ -9,7 +9,7 @@ import { OpenAI } from "llamaindex";
|
||||
|
||||
// chat api
|
||||
const response2 = await llm.chat([
|
||||
{ content: "Tell me a joke!", role: "user" },
|
||||
{ content: "Tell me a joke.", role: "user" },
|
||||
]);
|
||||
console.log(response2.message.content);
|
||||
})();
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
import { Portkey } from "llamaindex";
|
||||
|
||||
(async () => {
|
||||
const llms = [{}];
|
||||
const portkey = new Portkey({
|
||||
mode: "single",
|
||||
llms: [
|
||||
{
|
||||
provider: "anyscale",
|
||||
virtual_key: "anyscale-3b3c04",
|
||||
model: "meta-llama/Llama-2-13b-chat-hf",
|
||||
max_tokens: 2000,
|
||||
},
|
||||
],
|
||||
});
|
||||
const result = portkey.stream_chat([
|
||||
{ role: "system", content: "You are a helpful assistant." },
|
||||
{ role: "user", content: "Tell me a joke." },
|
||||
]);
|
||||
for await (const res of result) {
|
||||
process.stdout.write(res);
|
||||
}
|
||||
})();
|
||||
@@ -0,0 +1,37 @@
|
||||
import { execSync } from "child_process";
|
||||
import {
|
||||
PDFReader,
|
||||
serviceContextFromDefaults,
|
||||
storageContextFromDefaults,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
|
||||
const STORAGE_DIR = "./cache";
|
||||
|
||||
async function main() {
|
||||
// write the index to disk
|
||||
const serviceContext = serviceContextFromDefaults({});
|
||||
const storageContext = await storageContextFromDefaults({
|
||||
persistDir: `${STORAGE_DIR}`,
|
||||
});
|
||||
const reader = new PDFReader();
|
||||
const documents = await reader.loadData("data/brk-2022.pdf");
|
||||
await VectorStoreIndex.fromDocuments(documents, {
|
||||
storageContext,
|
||||
serviceContext,
|
||||
});
|
||||
console.log("wrote index to disk - now trying to read it");
|
||||
// make index dir read only
|
||||
execSync(`chmod -R 555 ${STORAGE_DIR}`);
|
||||
// reopen index
|
||||
const readOnlyStorageContext = await storageContextFromDefaults({
|
||||
persistDir: `${STORAGE_DIR}`,
|
||||
});
|
||||
await VectorStoreIndex.init({
|
||||
storageContext: readOnlyStorageContext,
|
||||
serviceContext,
|
||||
});
|
||||
console.log("read only index successfully opened");
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -2,7 +2,10 @@ import fs from "node:fs/promises";
|
||||
|
||||
import {
|
||||
Anthropic,
|
||||
anthropicTextQaPrompt,
|
||||
CompactAndRefine,
|
||||
Document,
|
||||
ResponseSynthesizer,
|
||||
serviceContextFromDefaults,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
@@ -18,12 +21,20 @@ async function main() {
|
||||
|
||||
// Split text and create embeddings. Store them in a VectorStoreIndex
|
||||
const serviceContext = serviceContextFromDefaults({ llm: new Anthropic() });
|
||||
|
||||
const responseSynthesizer = new ResponseSynthesizer({
|
||||
responseBuilder: new CompactAndRefine(
|
||||
serviceContext,
|
||||
anthropicTextQaPrompt,
|
||||
),
|
||||
});
|
||||
|
||||
const index = await VectorStoreIndex.fromDocuments([document], {
|
||||
serviceContext,
|
||||
});
|
||||
|
||||
// Query the index
|
||||
const queryEngine = index.asQueryEngine();
|
||||
const queryEngine = index.asQueryEngine({ responseSynthesizer });
|
||||
const response = await queryEngine.query(
|
||||
"What did the author do in college?",
|
||||
);
|
||||
|
||||
@@ -3,6 +3,7 @@ import {
|
||||
OpenAI,
|
||||
RetrieverQueryEngine,
|
||||
serviceContextFromDefaults,
|
||||
SimilarityPostprocessor,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
import essay from "./essay";
|
||||
@@ -21,8 +22,16 @@ async function main() {
|
||||
|
||||
const retriever = index.asRetriever();
|
||||
retriever.similarityTopK = 5;
|
||||
const nodePostprocessor = new SimilarityPostprocessor({
|
||||
similarityCutoff: 0.7,
|
||||
});
|
||||
// TODO: cannot pass responseSynthesizer into retriever query engine
|
||||
const queryEngine = new RetrieverQueryEngine(retriever);
|
||||
const queryEngine = new RetrieverQueryEngine(
|
||||
retriever,
|
||||
undefined,
|
||||
undefined,
|
||||
[nodePostprocessor],
|
||||
);
|
||||
|
||||
const response = await queryEngine.query(
|
||||
"What did the author do growing up?",
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
import {
|
||||
OpenAI,
|
||||
ResponseSynthesizer,
|
||||
RetrieverQueryEngine,
|
||||
serviceContextFromDefaults,
|
||||
TextNode,
|
||||
TreeSummarize,
|
||||
VectorIndexRetriever,
|
||||
VectorStore,
|
||||
VectorStoreIndex,
|
||||
VectorStoreQuery,
|
||||
VectorStoreQueryResult,
|
||||
} from "llamaindex";
|
||||
|
||||
import { Index, Pinecone, RecordMetadata } from "@pinecone-database/pinecone";
|
||||
|
||||
/**
|
||||
* Please do not use this class in production; it's only for demonstration purposes.
|
||||
*/
|
||||
class PineconeVectorStore<T extends RecordMetadata = RecordMetadata>
|
||||
implements VectorStore
|
||||
{
|
||||
storesText = true;
|
||||
isEmbeddingQuery = false;
|
||||
|
||||
indexName!: string;
|
||||
pineconeClient!: Pinecone;
|
||||
index!: Index<T>;
|
||||
|
||||
constructor({ indexName, client }: { indexName: string; client: Pinecone }) {
|
||||
this.indexName = indexName;
|
||||
this.pineconeClient = client;
|
||||
this.index = client.index<T>(indexName);
|
||||
}
|
||||
|
||||
client() {
|
||||
return this.pineconeClient;
|
||||
}
|
||||
|
||||
async query(
|
||||
query: VectorStoreQuery,
|
||||
kwargs?: any,
|
||||
): Promise<VectorStoreQueryResult> {
|
||||
let queryEmbedding: number[] = [];
|
||||
if (query.queryEmbedding) {
|
||||
if (typeof query.alpha === "number") {
|
||||
const alpha = query.alpha;
|
||||
queryEmbedding = query.queryEmbedding.map((v) => v * alpha);
|
||||
} else {
|
||||
queryEmbedding = query.queryEmbedding;
|
||||
}
|
||||
}
|
||||
|
||||
// Current LlamaIndexTS implementation only support exact match filter, so we use kwargs instead.
|
||||
const filter = kwargs?.filter || {};
|
||||
|
||||
const response = await this.index.query({
|
||||
filter,
|
||||
vector: queryEmbedding,
|
||||
topK: query.similarityTopK,
|
||||
includeValues: true,
|
||||
includeMetadata: true,
|
||||
});
|
||||
|
||||
console.log(
|
||||
`Numbers of vectors returned by Pinecone after preFilters are applied: ${
|
||||
response?.matches?.length || 0
|
||||
}.`,
|
||||
);
|
||||
|
||||
const topKIds: string[] = [];
|
||||
const topKNodes: TextNode[] = [];
|
||||
const topKScores: number[] = [];
|
||||
|
||||
const metadataToNode = (metadata?: T): Partial<TextNode> => {
|
||||
if (!metadata) {
|
||||
throw new Error("metadata is undefined.");
|
||||
}
|
||||
|
||||
const nodeContent = metadata["_node_content"];
|
||||
if (!nodeContent) {
|
||||
throw new Error("nodeContent is undefined.");
|
||||
}
|
||||
|
||||
if (typeof nodeContent !== "string") {
|
||||
throw new Error("nodeContent is not a string.");
|
||||
}
|
||||
|
||||
return JSON.parse(nodeContent);
|
||||
};
|
||||
|
||||
if (response.matches) {
|
||||
for (const match of response.matches) {
|
||||
const node = new TextNode({
|
||||
...metadataToNode(match.metadata),
|
||||
embedding: match.values,
|
||||
});
|
||||
|
||||
topKIds.push(match.id);
|
||||
topKNodes.push(node);
|
||||
topKScores.push(match.score ?? 0);
|
||||
}
|
||||
}
|
||||
|
||||
const result = {
|
||||
ids: topKIds,
|
||||
nodes: topKNodes,
|
||||
similarities: topKScores,
|
||||
};
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
add(): Promise<string[]> {
|
||||
return Promise.resolve([]);
|
||||
}
|
||||
|
||||
delete(): Promise<void> {
|
||||
throw new Error("Method `delete` not implemented.");
|
||||
}
|
||||
|
||||
persist(): Promise<void> {
|
||||
throw new Error("Method `persist` not implemented.");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The goal of this example is to show how to use Pinecone as a vector store
|
||||
* for LlamaIndexTS with(out) preFilters.
|
||||
*
|
||||
* It should not be used in production like that,
|
||||
* as you might want to find a proper PineconeVectorStore implementation.
|
||||
*/
|
||||
async function main() {
|
||||
process.env.PINECONE_API_KEY = "Your Pinecone API Key.";
|
||||
process.env.PINECONE_ENVIRONMENT = "Your Pinecone Environment.";
|
||||
process.env.PINECONE_PROJECT_ID = "Your Pinecone Project ID.";
|
||||
process.env.PINECONE_INDEX_NAME = "Your Pinecone Index Name.";
|
||||
process.env.OPENAI_API_KEY = "Your OpenAI API Key.";
|
||||
process.env.OPENAI_API_ORGANIZATION = "Your OpenAI API Organization.";
|
||||
|
||||
const getPineconeVectorStore = async () => {
|
||||
return new PineconeVectorStore({
|
||||
indexName: process.env.PINECONE_INDEX_NAME || "index-name",
|
||||
client: new Pinecone(),
|
||||
});
|
||||
};
|
||||
|
||||
const getServiceContext = () => {
|
||||
const openAI = new OpenAI({
|
||||
model: "gpt-4",
|
||||
apiKey: process.env.OPENAI_API_KEY,
|
||||
});
|
||||
|
||||
return serviceContextFromDefaults({
|
||||
llm: openAI,
|
||||
});
|
||||
};
|
||||
|
||||
const getQueryEngine = async (filter: unknown) => {
|
||||
const vectorStore = await getPineconeVectorStore();
|
||||
const serviceContext = getServiceContext();
|
||||
|
||||
const vectorStoreIndex = await VectorStoreIndex.fromVectorStore(
|
||||
vectorStore,
|
||||
serviceContext,
|
||||
);
|
||||
|
||||
const retriever = new VectorIndexRetriever({
|
||||
index: vectorStoreIndex,
|
||||
similarityTopK: 500,
|
||||
});
|
||||
|
||||
const responseSynthesizer = new ResponseSynthesizer({
|
||||
serviceContext,
|
||||
responseBuilder: new TreeSummarize(serviceContext),
|
||||
});
|
||||
|
||||
return new RetrieverQueryEngine(retriever, responseSynthesizer, {
|
||||
filter,
|
||||
});
|
||||
};
|
||||
|
||||
// whatever is a key from your metadata
|
||||
const queryEngine = await getQueryEngine({
|
||||
whatever: {
|
||||
$gte: 1,
|
||||
$lte: 100,
|
||||
},
|
||||
});
|
||||
|
||||
const response = await queryEngine.query("How many results do you have?");
|
||||
|
||||
console.log(response.toString());
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,15 @@
|
||||
import { OpenAI } from "llamaindex";
|
||||
|
||||
(async () => {
|
||||
const llm = new OpenAI({ model: "gpt-4-vision-preview", temperature: 0.1 });
|
||||
|
||||
// complete api
|
||||
const response1 = await llm.complete("How are you?");
|
||||
console.log(response1.message.content);
|
||||
|
||||
// chat api
|
||||
const response2 = await llm.chat([
|
||||
{ content: "Tell me a joke!", role: "user" },
|
||||
]);
|
||||
console.log(response2.message.content);
|
||||
})();
|
||||
+16
-13
@@ -3,7 +3,7 @@
|
||||
"scripts": {
|
||||
"build": "turbo run build",
|
||||
"dev": "turbo run dev",
|
||||
"format": "prettier --write \"**/*.{ts,tsx,md}\"",
|
||||
"format": "prettier --write \"**/*.{js,jsx,ts,tsx,md}\"",
|
||||
"lint": "turbo run lint",
|
||||
"prepare": "husky install",
|
||||
"test": "turbo run test",
|
||||
@@ -11,24 +11,27 @@
|
||||
"publish-snapshot": "turbo run build lint test && changeset version --snapshot && changeset publish"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@turbo/gen": "^1.10.13",
|
||||
"@types/jest": "^29.5.4",
|
||||
"eslint": "^7.32.0",
|
||||
"@changesets/cli": "^2.26.2",
|
||||
"@turbo/gen": "^1.10.16",
|
||||
"@types/jest": "^29.5.10",
|
||||
"eslint": "^8.54.0",
|
||||
"eslint-config-custom": "workspace:*",
|
||||
"husky": "^8.0.3",
|
||||
"jest": "^29.6.4",
|
||||
"prettier": "^3.0.3",
|
||||
"prettier-plugin-organize-imports": "^3.2.3",
|
||||
"jest": "^29.7.0",
|
||||
"lint-staged": "^15.1.0",
|
||||
"prettier": "^3.1.0",
|
||||
"prettier-plugin-organize-imports": "^3.2.4",
|
||||
"ts-jest": "^29.1.1",
|
||||
"turbo": "^1.10.13"
|
||||
},
|
||||
"packageManager": "pnpm@7.15.0",
|
||||
"dependencies": {
|
||||
"@changesets/cli": "^2.26.2"
|
||||
"turbo": "^1.10.16"
|
||||
},
|
||||
"packageManager": "pnpm@8.10.5+sha256.a4bd9bb7b48214bbfcd95f264bd75bb70d100e5d4b58808f5cd6ab40c6ac21c5",
|
||||
"pnpm": {
|
||||
"overrides": {
|
||||
"trim": "1.0.1"
|
||||
"trim": "1.0.1",
|
||||
"@babel/traverse": "7.23.2"
|
||||
}
|
||||
},
|
||||
"lint-staged": {
|
||||
"*.{js,jsx,ts,tsx,md}": "prettier --write"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,83 @@
|
||||
# llamaindex
|
||||
|
||||
## 0.0.36
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Support for Claude 2.1
|
||||
- Add AssemblyAI integration (thanks @Swimburger)
|
||||
- Use cryptoJS (thanks @marcusschiesser)
|
||||
- Add PGVectorStore (thanks @mtutty)
|
||||
- Add CLIP embeddings (thanks @marcusschiesser)
|
||||
- Add MongoDB support (thanks @marcusschiesser)
|
||||
|
||||
## 0.0.35
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 63f2108: Add multimodal support (thanks @marcusschiesser)
|
||||
|
||||
## 0.0.34
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 2a27e21: Add support for gpt-3.5-turbo-1106
|
||||
|
||||
## 0.0.33
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 5e2e92c: gpt-4-1106-preview and gpt-4-vision-preview from OpenAI dev day
|
||||
|
||||
## 0.0.32
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 90c0b83: Add HTMLReader (thanks @mtutty)
|
||||
- dfd22aa: Add observer/filter to the SimpleDirectoryReader (thanks @mtutty)
|
||||
|
||||
## 0.0.31
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 6c55b2d: Give HistoryChatEngine pluggable options (thanks @marcusschiesser)
|
||||
- 8aa8c65: Add SimilarityPostProcessor (thanks @TomPenguin)
|
||||
- 6c55b2d: Added LLMMetadata (thanks @marcusschiesser)
|
||||
|
||||
## 0.0.30
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 139abad: Streaming improvements including Anthropic (thanks @kkang2097)
|
||||
- 139abad: Portkey integration (Thank you @noble-varghese)
|
||||
- eb0e994: Add export for PromptHelper (thanks @zigamall)
|
||||
- eb0e994: Publish ESM module again
|
||||
- 139abad: Pinecone demo (thanks @Einsenhorn)
|
||||
|
||||
## 0.0.29
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- a52143b: Added DocxReader for Word documents (thanks @jayantasamaddar)
|
||||
- 1b7fd95: Updated OpenAI streaming (thanks @kkang2097)
|
||||
- 0db3f41: Migrated to Tiktoken lite, which hopefully fixes the Windows issue
|
||||
|
||||
## 0.0.28
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 96bb657: Typesafe metadata (thanks @TomPenguin)
|
||||
- 96bb657: MongoReader (thanks @kkang2097)
|
||||
- 837854d: Make OutputParser less strict and add tests (Thanks @kkang2097)
|
||||
|
||||
## 0.0.27
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 4a5591b: Chat History summarization (thanks @marcusschiesser)
|
||||
- 4a5591b: Notion database support (thanks @TomPenguin)
|
||||
- 4a5591b: KeywordIndex (thanks @swk777)
|
||||
|
||||
## 0.0.26
|
||||
|
||||
### Patch Changes
|
||||
|
||||
+31
-15
@@ -1,37 +1,53 @@
|
||||
{
|
||||
"name": "llamaindex",
|
||||
"version": "0.0.26",
|
||||
"version": "0.0.36",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"@anthropic-ai/sdk": "^0.6.2",
|
||||
"@notionhq/client": "^2.2.12",
|
||||
"@anthropic-ai/sdk": "^0.9.1",
|
||||
"@notionhq/client": "^2.2.13",
|
||||
"@xenova/transformers": "^2.8.0",
|
||||
"crypto-js": "^4.2.0",
|
||||
"js-tiktoken": "^1.0.8",
|
||||
"lodash": "^4.17.21",
|
||||
"mammoth": "^1.6.0",
|
||||
"md-utils-ts": "^2.0.0",
|
||||
"mongodb": "^6.3.0",
|
||||
"notion-md-crawler": "^0.0.2",
|
||||
"openai": "^4.3.1",
|
||||
"openai": "^4.19.1",
|
||||
"papaparse": "^5.4.1",
|
||||
"pdf-parse": "^1.1.1",
|
||||
"pg": "^8.11.3",
|
||||
"pgvector": "^0.1.5",
|
||||
"portkey-ai": "^0.1.16",
|
||||
"rake-modified": "^1.0.8",
|
||||
"replicate": "^0.16.1",
|
||||
"tiktoken-node": "^0.0.6",
|
||||
"uuid": "^9.0.0",
|
||||
"replicate": "^0.21.1",
|
||||
"string-strip-html": "^13.4.3",
|
||||
"uuid": "^9.0.1",
|
||||
"wink-nlp": "^1.14.3"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/lodash": "^4.14.197",
|
||||
"@types/node": "^18.17.12",
|
||||
"@types/papaparse": "^5.3.8",
|
||||
"@types/pdf-parse": "^1.1.1",
|
||||
"@types/uuid": "^9.0.3",
|
||||
"@types/crypto-js": "^4.2.1",
|
||||
"@types/lodash": "^4.14.202",
|
||||
"@types/node": "^18.18.12",
|
||||
"@types/papaparse": "^5.3.13",
|
||||
"@types/pdf-parse": "^1.1.4",
|
||||
"@types/pg": "^8.10.7",
|
||||
"@types/uuid": "^9.0.7",
|
||||
"node-stdlib-browser": "^1.2.0",
|
||||
"tsup": "^7.2.0"
|
||||
"tsup": "^7.2.0",
|
||||
"typescript": "^5.3.2"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=18.0.0"
|
||||
},
|
||||
"types": "./dist/index.d.ts",
|
||||
"main": "./dist/index.js",
|
||||
"module": "./dist/index.mjs",
|
||||
"repository": "run-llama/LlamaIndexTS",
|
||||
"scripts": {
|
||||
"lint": "eslint .",
|
||||
"test": "jest",
|
||||
"build": "tsup src/index.ts --format esm,cjs --dts"
|
||||
"build": "tsup src/index.ts --format esm,cjs --dts",
|
||||
"dev": "tsup src/index.ts --format esm,cjs --dts --watch"
|
||||
}
|
||||
}
|
||||
}
|
||||
+261
-46
@@ -1,8 +1,6 @@
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
import { Event } from "./callbacks/CallbackManager";
|
||||
import { ChatHistory, SimpleChatHistory } from "./ChatHistory";
|
||||
import { ChatMessage, LLM, OpenAI } from "./llm/LLM";
|
||||
import { TextNode } from "./Node";
|
||||
import { ChatHistory } from "./ChatHistory";
|
||||
import { NodeWithScore, TextNode } from "./Node";
|
||||
import {
|
||||
CondenseQuestionPrompt,
|
||||
ContextSystemPrompt,
|
||||
@@ -14,6 +12,9 @@ import { BaseQueryEngine } from "./QueryEngine";
|
||||
import { Response } from "./Response";
|
||||
import { BaseRetriever } from "./Retriever";
|
||||
import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext";
|
||||
import { Event } from "./callbacks/CallbackManager";
|
||||
import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor";
|
||||
import { ChatMessage, LLM, OpenAI } from "./llm/LLM";
|
||||
|
||||
/**
|
||||
* A ChatEngine is used to handle back and forth chats between the application and the LLM.
|
||||
@@ -23,8 +24,16 @@ export interface ChatEngine {
|
||||
* Send message along with the class's current chat history to the LLM.
|
||||
* @param message
|
||||
* @param chatHistory optional chat history if you want to customize the chat history
|
||||
* @param streaming optional streaming flag, which auto-sets the return value if True.
|
||||
*/
|
||||
chat(message: string, chatHistory?: ChatMessage[]): Promise<Response>;
|
||||
chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
|
||||
>(
|
||||
message: string,
|
||||
chatHistory?: ChatMessage[],
|
||||
streaming?: T,
|
||||
): Promise<R>;
|
||||
|
||||
/**
|
||||
* Resets the chat history so that it's empty.
|
||||
@@ -44,13 +53,45 @@ export class SimpleChatEngine implements ChatEngine {
|
||||
this.llm = init?.llm ?? new OpenAI();
|
||||
}
|
||||
|
||||
async chat(message: string, chatHistory?: ChatMessage[]): Promise<Response> {
|
||||
async chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
|
||||
>(message: string, chatHistory?: ChatMessage[], streaming?: T): Promise<R> {
|
||||
//Streaming option
|
||||
if (streaming) {
|
||||
return this.streamChat(message, chatHistory) as R;
|
||||
}
|
||||
|
||||
//Non-streaming option
|
||||
chatHistory = chatHistory ?? this.chatHistory;
|
||||
chatHistory.push({ content: message, role: "user" });
|
||||
const response = await this.llm.chat(chatHistory);
|
||||
const response = await this.llm.chat(chatHistory, undefined);
|
||||
chatHistory.push(response.message);
|
||||
this.chatHistory = chatHistory;
|
||||
return new Response(response.message.content);
|
||||
return new Response(response.message.content) as R;
|
||||
}
|
||||
|
||||
protected async *streamChat(
|
||||
message: string,
|
||||
chatHistory?: ChatMessage[],
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
chatHistory = chatHistory ?? this.chatHistory;
|
||||
chatHistory.push({ content: message, role: "user" });
|
||||
const response_generator = await this.llm.chat(
|
||||
chatHistory,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
|
||||
var accumulator: string = "";
|
||||
for await (const part of response_generator) {
|
||||
accumulator += part;
|
||||
yield part;
|
||||
}
|
||||
|
||||
chatHistory.push({ content: accumulator, role: "assistant" });
|
||||
this.chatHistory = chatHistory;
|
||||
return;
|
||||
}
|
||||
|
||||
reset() {
|
||||
@@ -99,10 +140,14 @@ export class CondenseQuestionChatEngine implements ChatEngine {
|
||||
);
|
||||
}
|
||||
|
||||
async chat(
|
||||
async chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
|
||||
>(
|
||||
message: string,
|
||||
chatHistory?: ChatMessage[] | undefined,
|
||||
): Promise<Response> {
|
||||
streaming?: T,
|
||||
): Promise<R> {
|
||||
chatHistory = chatHistory ?? this.chatHistory;
|
||||
|
||||
const condensedQuestion = (
|
||||
@@ -114,7 +159,7 @@ export class CondenseQuestionChatEngine implements ChatEngine {
|
||||
chatHistory.push({ content: message, role: "user" });
|
||||
chatHistory.push({ content: response.response, role: "assistant" });
|
||||
|
||||
return response;
|
||||
return response as R;
|
||||
}
|
||||
|
||||
reset() {
|
||||
@@ -122,57 +167,117 @@ export class CondenseQuestionChatEngine implements ChatEngine {
|
||||
}
|
||||
}
|
||||
|
||||
export interface Context {
|
||||
message: ChatMessage;
|
||||
nodes: NodeWithScore[];
|
||||
}
|
||||
|
||||
export interface ContextGenerator {
|
||||
generate(message: string, parentEvent?: Event): Promise<Context>;
|
||||
}
|
||||
|
||||
export class DefaultContextGenerator implements ContextGenerator {
|
||||
retriever: BaseRetriever;
|
||||
contextSystemPrompt: ContextSystemPrompt;
|
||||
nodePostprocessors: BaseNodePostprocessor[];
|
||||
|
||||
constructor(init: {
|
||||
retriever: BaseRetriever;
|
||||
contextSystemPrompt?: ContextSystemPrompt;
|
||||
nodePostprocessors?: BaseNodePostprocessor[];
|
||||
}) {
|
||||
this.retriever = init.retriever;
|
||||
this.contextSystemPrompt =
|
||||
init?.contextSystemPrompt ?? defaultContextSystemPrompt;
|
||||
this.nodePostprocessors = init.nodePostprocessors || [];
|
||||
}
|
||||
|
||||
private applyNodePostprocessors(nodes: NodeWithScore[]) {
|
||||
return this.nodePostprocessors.reduce(
|
||||
(nodes, nodePostprocessor) => nodePostprocessor.postprocessNodes(nodes),
|
||||
nodes,
|
||||
);
|
||||
}
|
||||
|
||||
async generate(message: string, parentEvent?: Event): Promise<Context> {
|
||||
if (!parentEvent) {
|
||||
parentEvent = {
|
||||
id: uuidv4(),
|
||||
type: "wrapper",
|
||||
tags: ["final"],
|
||||
};
|
||||
}
|
||||
const sourceNodesWithScore = await this.retriever.retrieve(
|
||||
message,
|
||||
parentEvent,
|
||||
);
|
||||
|
||||
const nodes = this.applyNodePostprocessors(sourceNodesWithScore);
|
||||
|
||||
return {
|
||||
message: {
|
||||
content: this.contextSystemPrompt({
|
||||
context: nodes.map((r) => (r.node as TextNode).text).join("\n\n"),
|
||||
}),
|
||||
role: "system",
|
||||
},
|
||||
nodes,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 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 preserved,
|
||||
* ideally allowing the appropriate context to be surfaced for each query.
|
||||
*/
|
||||
export class ContextChatEngine implements ChatEngine {
|
||||
retriever: BaseRetriever;
|
||||
chatModel: OpenAI;
|
||||
chatModel: LLM;
|
||||
chatHistory: ChatMessage[];
|
||||
contextSystemPrompt: ContextSystemPrompt;
|
||||
contextGenerator: ContextGenerator;
|
||||
|
||||
constructor(init: {
|
||||
retriever: BaseRetriever;
|
||||
chatModel?: OpenAI;
|
||||
chatModel?: LLM;
|
||||
chatHistory?: ChatMessage[];
|
||||
contextSystemPrompt?: ContextSystemPrompt;
|
||||
nodePostprocessors?: BaseNodePostprocessor[];
|
||||
}) {
|
||||
this.retriever = init.retriever;
|
||||
this.chatModel =
|
||||
init.chatModel ?? new OpenAI({ model: "gpt-3.5-turbo-16k" });
|
||||
this.chatHistory = init?.chatHistory ?? [];
|
||||
this.contextSystemPrompt =
|
||||
init?.contextSystemPrompt ?? defaultContextSystemPrompt;
|
||||
this.contextGenerator = new DefaultContextGenerator({
|
||||
retriever: init.retriever,
|
||||
contextSystemPrompt: init?.contextSystemPrompt,
|
||||
});
|
||||
}
|
||||
|
||||
async chat(message: string, chatHistory?: ChatMessage[] | undefined) {
|
||||
async chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
|
||||
>(
|
||||
message: string,
|
||||
chatHistory?: ChatMessage[] | undefined,
|
||||
streaming?: T,
|
||||
): Promise<R> {
|
||||
chatHistory = chatHistory ?? this.chatHistory;
|
||||
|
||||
//Streaming option
|
||||
if (streaming) {
|
||||
return this.streamChat(message, chatHistory) as R;
|
||||
}
|
||||
|
||||
const parentEvent: Event = {
|
||||
id: uuidv4(),
|
||||
type: "wrapper",
|
||||
tags: ["final"],
|
||||
};
|
||||
const sourceNodesWithScore = await this.retriever.retrieve(
|
||||
message,
|
||||
parentEvent,
|
||||
);
|
||||
|
||||
const systemMessage: ChatMessage = {
|
||||
content: this.contextSystemPrompt({
|
||||
context: sourceNodesWithScore
|
||||
.map((r) => (r.node as TextNode).text)
|
||||
.join("\n\n"),
|
||||
}),
|
||||
role: "system",
|
||||
};
|
||||
const context = await this.contextGenerator.generate(message, parentEvent);
|
||||
|
||||
chatHistory.push({ content: message, role: "user" });
|
||||
|
||||
const response = await this.chatModel.chat(
|
||||
[systemMessage, ...chatHistory],
|
||||
[context.message, ...chatHistory],
|
||||
parentEvent,
|
||||
);
|
||||
chatHistory.push(response.message);
|
||||
@@ -181,8 +286,41 @@ export class ContextChatEngine implements ChatEngine {
|
||||
|
||||
return new Response(
|
||||
response.message.content,
|
||||
sourceNodesWithScore.map((r) => r.node),
|
||||
context.nodes.map((r) => r.node),
|
||||
) as R;
|
||||
}
|
||||
|
||||
protected async *streamChat(
|
||||
message: string,
|
||||
chatHistory?: ChatMessage[] | undefined,
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
chatHistory = chatHistory ?? this.chatHistory;
|
||||
|
||||
const parentEvent: Event = {
|
||||
id: uuidv4(),
|
||||
type: "wrapper",
|
||||
tags: ["final"],
|
||||
};
|
||||
const context = await this.contextGenerator.generate(message, parentEvent);
|
||||
|
||||
chatHistory.push({ content: message, role: "user" });
|
||||
|
||||
const response_stream = await this.chatModel.chat(
|
||||
[context.message, ...chatHistory],
|
||||
parentEvent,
|
||||
true,
|
||||
);
|
||||
var accumulator: string = "";
|
||||
for await (const part of response_stream) {
|
||||
accumulator += part;
|
||||
yield part;
|
||||
}
|
||||
|
||||
chatHistory.push({ content: accumulator, role: "assistant" });
|
||||
|
||||
this.chatHistory = chatHistory;
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
reset() {
|
||||
@@ -190,27 +328,104 @@ export class ContextChatEngine implements ChatEngine {
|
||||
}
|
||||
}
|
||||
|
||||
export interface MessageContentDetail {
|
||||
type: "text" | "image_url";
|
||||
text: string;
|
||||
image_url: { url: string };
|
||||
}
|
||||
|
||||
/**
|
||||
* HistoryChatEngine is a ChatEngine that uses a ChatHistory to keep track of the chat history. This is an example with the same behavior as SimpleChatEngine
|
||||
* TODO: generally use the ChatHistory instead of ChatMessage[] - breaking change
|
||||
* Extended type for the content of a message that allows for multi-modal messages.
|
||||
*/
|
||||
export class HistoryChatEngine implements ChatEngine {
|
||||
chatHistory: ChatHistory;
|
||||
export type MessageContent = string | MessageContentDetail[];
|
||||
|
||||
/**
|
||||
* HistoryChatEngine is a ChatEngine that uses a `ChatHistory` object
|
||||
* to keeps track of chat's message history.
|
||||
* A `ChatHistory` object is passed as a parameter for each call to the `chat` method,
|
||||
* so the state of the chat engine is preserved between calls.
|
||||
* Optionally, a `ContextGenerator` can be used to generate an additional context for each call to `chat`.
|
||||
*/
|
||||
export class HistoryChatEngine {
|
||||
llm: LLM;
|
||||
contextGenerator?: ContextGenerator;
|
||||
|
||||
constructor(init?: Partial<HistoryChatEngine>) {
|
||||
this.chatHistory = init?.chatHistory ?? new SimpleChatHistory();
|
||||
this.llm = init?.llm ?? new OpenAI();
|
||||
this.contextGenerator = init?.contextGenerator;
|
||||
}
|
||||
|
||||
async chat(message: string): Promise<Response> {
|
||||
this.chatHistory.addMessage({ content: message, role: "user" });
|
||||
const response = await this.llm.chat(this.chatHistory.messages);
|
||||
this.chatHistory.addMessage(response.message);
|
||||
return new Response(response.message.content);
|
||||
async chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
|
||||
>(
|
||||
message: MessageContent,
|
||||
chatHistory: ChatHistory,
|
||||
streaming?: T,
|
||||
): Promise<R> {
|
||||
//Streaming option
|
||||
if (streaming) {
|
||||
return this.streamChat(message, chatHistory) as R;
|
||||
}
|
||||
const requestMessages = await this.prepareRequestMessages(
|
||||
message,
|
||||
chatHistory,
|
||||
);
|
||||
const response = await this.llm.chat(requestMessages);
|
||||
chatHistory.addMessage(response.message);
|
||||
return new Response(response.message.content) as R;
|
||||
}
|
||||
|
||||
reset() {
|
||||
this.chatHistory.reset();
|
||||
protected async *streamChat(
|
||||
message: MessageContent,
|
||||
chatHistory: ChatHistory,
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
const requestMessages = await this.prepareRequestMessages(
|
||||
message,
|
||||
chatHistory,
|
||||
);
|
||||
const response_stream = await this.llm.chat(
|
||||
requestMessages,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
|
||||
var accumulator = "";
|
||||
for await (const part of response_stream) {
|
||||
accumulator += part;
|
||||
yield part;
|
||||
}
|
||||
chatHistory.addMessage({
|
||||
content: accumulator,
|
||||
role: "assistant",
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
private async prepareRequestMessages(
|
||||
message: MessageContent,
|
||||
chatHistory: ChatHistory,
|
||||
) {
|
||||
chatHistory.addMessage({
|
||||
content: message,
|
||||
role: "user",
|
||||
});
|
||||
let requestMessages;
|
||||
let context;
|
||||
if (this.contextGenerator) {
|
||||
if (Array.isArray(message)) {
|
||||
// message is of type MessageContentDetail[] - retrieve just the text parts and concatenate them
|
||||
// so we can pass them to the context generator
|
||||
message = (message as MessageContentDetail[])
|
||||
.filter((c) => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("\n\n");
|
||||
}
|
||||
context = await this.contextGenerator.generate(message);
|
||||
}
|
||||
requestMessages = await chatHistory.requestMessages(
|
||||
context ? [context.message] : undefined,
|
||||
);
|
||||
return requestMessages;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { ChatMessage, LLM, OpenAI } from "./llm/LLM";
|
||||
import { ChatMessage, LLM, MessageType, OpenAI } from "./llm/LLM";
|
||||
import {
|
||||
defaultSummaryPrompt,
|
||||
messagesToHistoryStr,
|
||||
@@ -14,60 +14,187 @@ export interface ChatHistory {
|
||||
* Adds a message to the chat history.
|
||||
* @param message
|
||||
*/
|
||||
addMessage(message: ChatMessage): Promise<void>;
|
||||
addMessage(message: ChatMessage): void;
|
||||
|
||||
/**
|
||||
* Returns the messages that should be used as input to the LLM.
|
||||
*/
|
||||
requestMessages(transientMessages?: ChatMessage[]): Promise<ChatMessage[]>;
|
||||
|
||||
/**
|
||||
* Resets the chat history so that it's empty.
|
||||
*/
|
||||
reset(): void;
|
||||
|
||||
/**
|
||||
* Returns the new messages since the last call to this function (or since calling the constructor)
|
||||
*/
|
||||
newMessages(): ChatMessage[];
|
||||
}
|
||||
|
||||
export class SimpleChatHistory implements ChatHistory {
|
||||
messages: ChatMessage[];
|
||||
private messagesBefore: number;
|
||||
|
||||
constructor(init?: Partial<SimpleChatHistory>) {
|
||||
this.messages = init?.messages ?? [];
|
||||
this.messagesBefore = this.messages.length;
|
||||
}
|
||||
|
||||
async addMessage(message: ChatMessage) {
|
||||
addMessage(message: ChatMessage) {
|
||||
this.messages.push(message);
|
||||
}
|
||||
|
||||
async requestMessages(transientMessages?: ChatMessage[]) {
|
||||
return [...(transientMessages ?? []), ...this.messages];
|
||||
}
|
||||
|
||||
reset() {
|
||||
this.messages = [];
|
||||
}
|
||||
|
||||
newMessages() {
|
||||
const newMessages = this.messages.slice(this.messagesBefore);
|
||||
this.messagesBefore = this.messages.length;
|
||||
return newMessages;
|
||||
}
|
||||
}
|
||||
|
||||
export class SummaryChatHistory implements ChatHistory {
|
||||
tokensToSummarize: number;
|
||||
messages: ChatMessage[];
|
||||
summaryPrompt: SummaryPrompt;
|
||||
llm: LLM;
|
||||
private messagesBefore: number;
|
||||
|
||||
constructor(init?: Partial<SummaryChatHistory>) {
|
||||
this.messages = init?.messages ?? [];
|
||||
this.messagesBefore = this.messages.length;
|
||||
this.summaryPrompt = init?.summaryPrompt ?? defaultSummaryPrompt;
|
||||
this.llm = init?.llm ?? new OpenAI();
|
||||
if (!this.llm.metadata.maxTokens) {
|
||||
throw new Error(
|
||||
"LLM maxTokens is not set. Needed so the summarizer ensures the context window size of the LLM.",
|
||||
);
|
||||
}
|
||||
this.tokensToSummarize =
|
||||
this.llm.metadata.contextWindow - this.llm.metadata.maxTokens;
|
||||
}
|
||||
|
||||
private async summarize() {
|
||||
const chatHistoryStr = messagesToHistoryStr(this.messages);
|
||||
private async summarize(): Promise<ChatMessage> {
|
||||
// get the conversation messages to create summary
|
||||
const messagesToSummarize = this.calcConversationMessages();
|
||||
|
||||
const response = await this.llm.complete(
|
||||
this.summaryPrompt({ context: chatHistoryStr }),
|
||||
);
|
||||
let promptMessages;
|
||||
do {
|
||||
promptMessages = [
|
||||
{
|
||||
content: this.summaryPrompt({
|
||||
context: messagesToHistoryStr(messagesToSummarize),
|
||||
}),
|
||||
role: "user" as MessageType,
|
||||
},
|
||||
];
|
||||
// remove oldest message until the chat history is short enough for the context window
|
||||
messagesToSummarize.shift();
|
||||
} while (this.llm.tokens(promptMessages) > this.tokensToSummarize);
|
||||
|
||||
this.messages = [{ content: response.message.content, role: "system" }];
|
||||
const response = await this.llm.chat(promptMessages);
|
||||
return { content: response.message.content, role: "memory" };
|
||||
}
|
||||
|
||||
async addMessage(message: ChatMessage) {
|
||||
// TODO: check if summarization is necessary
|
||||
// TBD what are good conditions, e.g. depending on the context length of the LLM?
|
||||
// for now we just have a dummy implementation at always summarizes the messages
|
||||
await this.summarize();
|
||||
addMessage(message: ChatMessage) {
|
||||
this.messages.push(message);
|
||||
}
|
||||
|
||||
// Find last summary message
|
||||
private getLastSummaryIndex(): number | null {
|
||||
const reversedMessages = this.messages.slice().reverse();
|
||||
const index = reversedMessages.findIndex(
|
||||
(message) => message.role === "memory",
|
||||
);
|
||||
if (index === -1) {
|
||||
return null;
|
||||
}
|
||||
return this.messages.length - 1 - index;
|
||||
}
|
||||
|
||||
private get systemMessages() {
|
||||
// get array of all system messages
|
||||
return this.messages.filter((message) => message.role === "system");
|
||||
}
|
||||
|
||||
private get nonSystemMessages() {
|
||||
// get array of all non-system messages
|
||||
return this.messages.filter((message) => message.role !== "system");
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculates the messages that describe the conversation so far.
|
||||
* If there's no memory, all non-system messages are used.
|
||||
* If there's a memory, uses all messages after the last summary message.
|
||||
*/
|
||||
private calcConversationMessages(transformSummary?: boolean): ChatMessage[] {
|
||||
const lastSummaryIndex = this.getLastSummaryIndex();
|
||||
if (!lastSummaryIndex) {
|
||||
// there's no memory, so just use all non-system messages
|
||||
return this.nonSystemMessages;
|
||||
} else {
|
||||
// there's a memory, so use all messages after the last summary message
|
||||
// and convert summary message so it can be send to the LLM
|
||||
const summaryMessage: ChatMessage = transformSummary
|
||||
? {
|
||||
content: `Summary of the conversation so far: ${this.messages[lastSummaryIndex].content}`,
|
||||
role: "system",
|
||||
}
|
||||
: this.messages[lastSummaryIndex];
|
||||
return [summaryMessage, ...this.messages.slice(lastSummaryIndex + 1)];
|
||||
}
|
||||
}
|
||||
|
||||
private calcCurrentRequestMessages(transientMessages?: ChatMessage[]) {
|
||||
// TODO: check order: currently, we're sending:
|
||||
// system messages first, then transient messages and then the messages that describe the conversation so far
|
||||
return [
|
||||
...this.systemMessages,
|
||||
...(transientMessages ? transientMessages : []),
|
||||
...this.calcConversationMessages(true),
|
||||
];
|
||||
}
|
||||
|
||||
async requestMessages(transientMessages?: ChatMessage[]) {
|
||||
const requestMessages = this.calcCurrentRequestMessages(transientMessages);
|
||||
|
||||
// get tokens of current request messages and the transient messages
|
||||
const tokens = this.llm.tokens(requestMessages);
|
||||
if (tokens > this.tokensToSummarize) {
|
||||
// if there are too many tokens for the next request, call summarize
|
||||
const memoryMessage = await this.summarize();
|
||||
const lastMessage = this.messages.at(-1);
|
||||
if (lastMessage && lastMessage.role === "user") {
|
||||
// if last message is a user message, ensure that it's sent after the new memory message
|
||||
this.messages.pop();
|
||||
this.messages.push(memoryMessage);
|
||||
this.messages.push(lastMessage);
|
||||
} else {
|
||||
// otherwise just add the memory message
|
||||
this.messages.push(memoryMessage);
|
||||
}
|
||||
// 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(transientMessages);
|
||||
}
|
||||
return requestMessages;
|
||||
}
|
||||
|
||||
reset() {
|
||||
this.messages = [];
|
||||
}
|
||||
|
||||
newMessages() {
|
||||
const newMessages = this.messages.slice(this.messagesBefore);
|
||||
this.messagesBefore = this.messages.length;
|
||||
return newMessages;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,28 +1,54 @@
|
||||
import { encodingForModel } from "js-tiktoken";
|
||||
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
import { Event, EventTag, EventType } from "./callbacks/CallbackManager";
|
||||
|
||||
export enum Tokenizers {
|
||||
CL100K_BASE = "cl100k_base",
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper class singleton
|
||||
*/
|
||||
class GlobalsHelper {
|
||||
defaultTokenizer: {
|
||||
encode: (text: string) => number[];
|
||||
decode: (tokens: number[]) => string;
|
||||
encode: (text: string) => Uint32Array;
|
||||
decode: (tokens: Uint32Array) => string;
|
||||
} | null = null;
|
||||
|
||||
tokenizer() {
|
||||
private initDefaultTokenizer() {
|
||||
const encoding = encodingForModel("text-embedding-ada-002"); // cl100k_base
|
||||
|
||||
this.defaultTokenizer = {
|
||||
encode: (text: string) => {
|
||||
return new Uint32Array(encoding.encode(text));
|
||||
},
|
||||
decode: (tokens: Uint32Array) => {
|
||||
const numberArray = Array.from(tokens);
|
||||
const text = encoding.decode(numberArray);
|
||||
const uint8Array = new TextEncoder().encode(text);
|
||||
return new TextDecoder().decode(uint8Array);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
tokenizer(encoding?: string) {
|
||||
if (encoding && encoding !== Tokenizers.CL100K_BASE) {
|
||||
throw new Error(`Tokenizer encoding ${encoding} not yet supported`);
|
||||
}
|
||||
if (!this.defaultTokenizer) {
|
||||
const tiktoken = require("tiktoken-node");
|
||||
this.defaultTokenizer = tiktoken.getEncoding("gpt2");
|
||||
this.initDefaultTokenizer();
|
||||
}
|
||||
|
||||
return this.defaultTokenizer!.encode.bind(this.defaultTokenizer);
|
||||
}
|
||||
|
||||
tokenizerDecoder() {
|
||||
tokenizerDecoder(encoding?: string) {
|
||||
if (encoding && encoding !== Tokenizers.CL100K_BASE) {
|
||||
throw new Error(`Tokenizer encoding ${encoding} not yet supported`);
|
||||
}
|
||||
if (!this.defaultTokenizer) {
|
||||
const tiktoken = require("tiktoken-node");
|
||||
this.defaultTokenizer = tiktoken.getEncoding("gpt2");
|
||||
this.initDefaultTokenizer();
|
||||
}
|
||||
|
||||
return this.defaultTokenizer!.decode.bind(this.defaultTokenizer);
|
||||
|
||||
+33
-28
@@ -1,4 +1,4 @@
|
||||
import crypto from "crypto"; // TODO Node dependency
|
||||
import CryptoJS from "crypto-js";
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
|
||||
export enum NodeRelationship {
|
||||
@@ -23,19 +23,23 @@ export enum MetadataMode {
|
||||
NONE = "NONE",
|
||||
}
|
||||
|
||||
export interface RelatedNodeInfo {
|
||||
export type Metadata = Record<string, any>;
|
||||
|
||||
export interface RelatedNodeInfo<T extends Metadata = Metadata> {
|
||||
nodeId: string;
|
||||
nodeType?: ObjectType;
|
||||
metadata: Record<string, any>;
|
||||
metadata: T;
|
||||
hash?: string;
|
||||
}
|
||||
|
||||
export type RelatedNodeType = RelatedNodeInfo | RelatedNodeInfo[];
|
||||
export type RelatedNodeType<T extends Metadata = Metadata> =
|
||||
| RelatedNodeInfo<T>
|
||||
| RelatedNodeInfo<T>[];
|
||||
|
||||
/**
|
||||
* Generic abstract class for retrievable nodes
|
||||
*/
|
||||
export abstract class BaseNode {
|
||||
export abstract class BaseNode<T extends Metadata = Metadata> {
|
||||
/**
|
||||
* The unique ID of the Node/Document. The trailing underscore is here
|
||||
* to avoid collisions with the id keyword in Python.
|
||||
@@ -46,13 +50,13 @@ export abstract class BaseNode {
|
||||
embedding?: number[];
|
||||
|
||||
// Metadata fields
|
||||
metadata: Record<string, any> = {};
|
||||
metadata: T = {} as T;
|
||||
excludedEmbedMetadataKeys: string[] = [];
|
||||
excludedLlmMetadataKeys: string[] = [];
|
||||
relationships: Partial<Record<NodeRelationship, RelatedNodeType>> = {};
|
||||
relationships: Partial<Record<NodeRelationship, RelatedNodeType<T>>> = {};
|
||||
hash: string = "";
|
||||
|
||||
constructor(init?: Partial<BaseNode>) {
|
||||
constructor(init?: Partial<BaseNode<T>>) {
|
||||
Object.assign(this, init);
|
||||
}
|
||||
|
||||
@@ -62,7 +66,7 @@ export abstract class BaseNode {
|
||||
abstract getMetadataStr(metadataMode: MetadataMode): string;
|
||||
abstract setContent(value: any): void;
|
||||
|
||||
get sourceNode(): RelatedNodeInfo | undefined {
|
||||
get sourceNode(): RelatedNodeInfo<T> | undefined {
|
||||
const relationship = this.relationships[NodeRelationship.SOURCE];
|
||||
|
||||
if (Array.isArray(relationship)) {
|
||||
@@ -72,7 +76,7 @@ export abstract class BaseNode {
|
||||
return relationship;
|
||||
}
|
||||
|
||||
get prevNode(): RelatedNodeInfo | undefined {
|
||||
get prevNode(): RelatedNodeInfo<T> | undefined {
|
||||
const relationship = this.relationships[NodeRelationship.PREVIOUS];
|
||||
|
||||
if (Array.isArray(relationship)) {
|
||||
@@ -84,7 +88,7 @@ export abstract class BaseNode {
|
||||
return relationship;
|
||||
}
|
||||
|
||||
get nextNode(): RelatedNodeInfo | undefined {
|
||||
get nextNode(): RelatedNodeInfo<T> | undefined {
|
||||
const relationship = this.relationships[NodeRelationship.NEXT];
|
||||
|
||||
if (Array.isArray(relationship)) {
|
||||
@@ -94,7 +98,7 @@ export abstract class BaseNode {
|
||||
return relationship;
|
||||
}
|
||||
|
||||
get parentNode(): RelatedNodeInfo | undefined {
|
||||
get parentNode(): RelatedNodeInfo<T> | undefined {
|
||||
const relationship = this.relationships[NodeRelationship.PARENT];
|
||||
|
||||
if (Array.isArray(relationship)) {
|
||||
@@ -104,7 +108,7 @@ export abstract class BaseNode {
|
||||
return relationship;
|
||||
}
|
||||
|
||||
get childNodes(): RelatedNodeInfo[] | undefined {
|
||||
get childNodes(): RelatedNodeInfo<T>[] | undefined {
|
||||
const relationship = this.relationships[NodeRelationship.CHILD];
|
||||
|
||||
if (!Array.isArray(relationship)) {
|
||||
@@ -126,7 +130,7 @@ export abstract class BaseNode {
|
||||
return this.embedding;
|
||||
}
|
||||
|
||||
asRelatedNodeInfo(): RelatedNodeInfo {
|
||||
asRelatedNodeInfo(): RelatedNodeInfo<T> {
|
||||
return {
|
||||
nodeId: this.id_,
|
||||
metadata: this.metadata,
|
||||
@@ -146,7 +150,7 @@ export abstract class BaseNode {
|
||||
/**
|
||||
* TextNode is the default node type for text. Most common node type in LlamaIndex.TS
|
||||
*/
|
||||
export class TextNode extends BaseNode {
|
||||
export class TextNode<T extends Metadata = Metadata> extends BaseNode<T> {
|
||||
text: string = "";
|
||||
startCharIdx?: number;
|
||||
endCharIdx?: number;
|
||||
@@ -154,7 +158,7 @@ export class TextNode extends BaseNode {
|
||||
// metadataTemplate: NOTE write your own formatter if needed
|
||||
metadataSeparator: string = "\n";
|
||||
|
||||
constructor(init?: Partial<TextNode>) {
|
||||
constructor(init?: Partial<TextNode<T>>) {
|
||||
super(init);
|
||||
Object.assign(this, init);
|
||||
|
||||
@@ -171,13 +175,13 @@ export class TextNode extends BaseNode {
|
||||
* @returns
|
||||
*/
|
||||
generateHash() {
|
||||
const hashFunction = crypto.createHash("sha256");
|
||||
const hashFunction = CryptoJS.algo.SHA256.create();
|
||||
hashFunction.update(`type=${this.getType()}`);
|
||||
hashFunction.update(
|
||||
`startCharIdx=${this.startCharIdx} endCharIdx=${this.endCharIdx}`,
|
||||
);
|
||||
hashFunction.update(this.getContent(MetadataMode.ALL));
|
||||
return hashFunction.digest("base64");
|
||||
return hashFunction.finalize().toString(CryptoJS.enc.Base64);
|
||||
}
|
||||
|
||||
getType(): ObjectType {
|
||||
@@ -233,10 +237,10 @@ export class TextNode extends BaseNode {
|
||||
// }
|
||||
// }
|
||||
|
||||
export class IndexNode extends TextNode {
|
||||
export class IndexNode<T extends Metadata = Metadata> extends TextNode<T> {
|
||||
indexId: string = "";
|
||||
|
||||
constructor(init?: Partial<IndexNode>) {
|
||||
constructor(init?: Partial<IndexNode<T>>) {
|
||||
super(init);
|
||||
Object.assign(this, init);
|
||||
|
||||
@@ -253,8 +257,8 @@ export class IndexNode extends TextNode {
|
||||
/**
|
||||
* A document is just a special text node with a docId.
|
||||
*/
|
||||
export class Document extends TextNode {
|
||||
constructor(init?: Partial<Document>) {
|
||||
export class Document<T extends Metadata = Metadata> extends TextNode<T> {
|
||||
constructor(init?: Partial<Document<T>>) {
|
||||
super(init);
|
||||
Object.assign(this, init);
|
||||
|
||||
@@ -268,12 +272,13 @@ export class Document extends TextNode {
|
||||
}
|
||||
}
|
||||
|
||||
export function jsonToNode(json: any) {
|
||||
if (!json.type) {
|
||||
export function jsonToNode(json: any, type?: ObjectType) {
|
||||
if (!json.type && !type) {
|
||||
throw new Error("Node type not found");
|
||||
}
|
||||
const nodeType = type || json.type;
|
||||
|
||||
switch (json.type) {
|
||||
switch (nodeType) {
|
||||
case ObjectType.TEXT:
|
||||
return new TextNode(json);
|
||||
case ObjectType.INDEX:
|
||||
@@ -281,7 +286,7 @@ export function jsonToNode(json: any) {
|
||||
case ObjectType.DOCUMENT:
|
||||
return new Document(json);
|
||||
default:
|
||||
throw new Error(`Invalid node type: ${json.type}`);
|
||||
throw new Error(`Invalid node type: ${nodeType}`);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -292,7 +297,7 @@ export function jsonToNode(json: any) {
|
||||
/**
|
||||
* A node with a similarity score
|
||||
*/
|
||||
export interface NodeWithScore {
|
||||
node: BaseNode;
|
||||
export interface NodeWithScore<T extends Metadata = Metadata> {
|
||||
node: BaseNode<T>;
|
||||
score?: number;
|
||||
}
|
||||
|
||||
@@ -53,30 +53,31 @@ class OutputParserError extends Error {
|
||||
* @param text A markdown block with JSON
|
||||
* @returns parsed JSON object
|
||||
*/
|
||||
function parseJsonMarkdown(text: string) {
|
||||
export function parseJsonMarkdown(text: string) {
|
||||
text = text.trim();
|
||||
|
||||
const beginDelimiter = "```json";
|
||||
const endDelimiter = "```";
|
||||
const left_square = text.indexOf("[");
|
||||
const left_brace = text.indexOf("{");
|
||||
|
||||
const beginIndex = text.indexOf(beginDelimiter);
|
||||
const endIndex = text.indexOf(
|
||||
endDelimiter,
|
||||
beginIndex + beginDelimiter.length,
|
||||
);
|
||||
if (beginIndex === -1 || endIndex === -1) {
|
||||
throw new OutputParserError("Not a json markdown", { output: text });
|
||||
var left: number;
|
||||
var right: number;
|
||||
if (left_square < left_brace && left_square != -1) {
|
||||
left = left_square;
|
||||
right = text.lastIndexOf("]");
|
||||
} else {
|
||||
left = left_brace;
|
||||
right = text.lastIndexOf("}");
|
||||
}
|
||||
|
||||
const jsonText = text.substring(beginIndex + beginDelimiter.length, endIndex);
|
||||
|
||||
const jsonText = text.substring(left, right + 1);
|
||||
try {
|
||||
//Single JSON object case
|
||||
if (left_square === -1) {
|
||||
return [JSON.parse(jsonText)];
|
||||
}
|
||||
//Multiple JSON object case.
|
||||
return JSON.parse(jsonText);
|
||||
} catch (e) {
|
||||
throw new OutputParserError("Not a valid json", {
|
||||
cause: e as Error,
|
||||
output: text,
|
||||
});
|
||||
throw new OutputParserError("Not a json markdown", { output: text });
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -36,6 +36,15 @@ Answer:`;
|
||||
|
||||
export type TextQaPrompt = typeof defaultTextQaPrompt;
|
||||
|
||||
export const anthropicTextQaPrompt = ({ context = "", query = "" }) => {
|
||||
return `Context information:
|
||||
<context>
|
||||
${context}
|
||||
</context>
|
||||
Given the context information and not prior knowledge, answer the query.
|
||||
Query: ${query}`;
|
||||
};
|
||||
|
||||
/*
|
||||
DEFAULT_SUMMARY_PROMPT_TMPL = (
|
||||
"Write a summary of the following. Try to use only the "
|
||||
|
||||
@@ -34,7 +34,7 @@ export class PromptHelper {
|
||||
numOutput = DEFAULT_NUM_OUTPUTS;
|
||||
chunkOverlapRatio = DEFAULT_CHUNK_OVERLAP_RATIO;
|
||||
chunkSizeLimit?: number;
|
||||
tokenizer: (text: string) => number[];
|
||||
tokenizer: (text: string) => Uint32Array;
|
||||
separator = " ";
|
||||
|
||||
constructor(
|
||||
@@ -42,7 +42,7 @@ export class PromptHelper {
|
||||
numOutput = DEFAULT_NUM_OUTPUTS,
|
||||
chunkOverlapRatio = DEFAULT_CHUNK_OVERLAP_RATIO,
|
||||
chunkSizeLimit?: number,
|
||||
tokenizer?: (text: string) => number[],
|
||||
tokenizer?: (text: string) => Uint32Array,
|
||||
separator = " ",
|
||||
) {
|
||||
this.contextWindow = contextWindow;
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
import { Event } from "./callbacks/CallbackManager";
|
||||
import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor";
|
||||
import { NodeWithScore, TextNode } from "./Node";
|
||||
import {
|
||||
BaseQuestionGenerator,
|
||||
@@ -10,7 +12,6 @@ import { CompactAndRefine, ResponseSynthesizer } from "./ResponseSynthesizer";
|
||||
import { BaseRetriever } from "./Retriever";
|
||||
import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext";
|
||||
import { QueryEngineTool, ToolMetadata } from "./Tool";
|
||||
import { Event } from "./callbacks/CallbackManager";
|
||||
|
||||
/**
|
||||
* A query engine is a question answerer that can use one or more steps.
|
||||
@@ -30,16 +31,39 @@ export interface BaseQueryEngine {
|
||||
export class RetrieverQueryEngine implements BaseQueryEngine {
|
||||
retriever: BaseRetriever;
|
||||
responseSynthesizer: ResponseSynthesizer;
|
||||
nodePostprocessors: BaseNodePostprocessor[];
|
||||
preFilters?: unknown;
|
||||
|
||||
constructor(
|
||||
retriever: BaseRetriever,
|
||||
responseSynthesizer?: ResponseSynthesizer,
|
||||
preFilters?: unknown,
|
||||
nodePostprocessors?: BaseNodePostprocessor[],
|
||||
) {
|
||||
this.retriever = retriever;
|
||||
const serviceContext: ServiceContext | undefined =
|
||||
this.retriever.getServiceContext();
|
||||
this.responseSynthesizer =
|
||||
responseSynthesizer || new ResponseSynthesizer({ serviceContext });
|
||||
this.preFilters = preFilters;
|
||||
this.nodePostprocessors = nodePostprocessors || [];
|
||||
}
|
||||
|
||||
private applyNodePostprocessors(nodes: NodeWithScore[]) {
|
||||
return this.nodePostprocessors.reduce(
|
||||
(nodes, nodePostprocessor) => nodePostprocessor.postprocessNodes(nodes),
|
||||
nodes,
|
||||
);
|
||||
}
|
||||
|
||||
private async retrieve(query: string, parentEvent: Event) {
|
||||
const nodes = await this.retriever.retrieve(
|
||||
query,
|
||||
parentEvent,
|
||||
this.preFilters,
|
||||
);
|
||||
|
||||
return this.applyNodePostprocessors(nodes);
|
||||
}
|
||||
|
||||
async query(query: string, parentEvent?: Event) {
|
||||
@@ -48,7 +72,7 @@ export class RetrieverQueryEngine implements BaseQueryEngine {
|
||||
type: "wrapper",
|
||||
tags: ["final"],
|
||||
};
|
||||
const nodes = await this.retriever.retrieve(query, _parentEvent);
|
||||
const nodes = await this.retrieve(query, _parentEvent);
|
||||
return this.responseSynthesizer.synthesize(query, nodes, _parentEvent);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,18 +1,18 @@
|
||||
import { Event } from "./callbacks/CallbackManager";
|
||||
import { LLM } from "./llm/LLM";
|
||||
import { MetadataMode, NodeWithScore } from "./Node";
|
||||
import {
|
||||
defaultRefinePrompt,
|
||||
defaultTextQaPrompt,
|
||||
defaultTreeSummarizePrompt,
|
||||
RefinePrompt,
|
||||
SimplePrompt,
|
||||
TextQaPrompt,
|
||||
TreeSummarizePrompt,
|
||||
defaultRefinePrompt,
|
||||
defaultTextQaPrompt,
|
||||
defaultTreeSummarizePrompt,
|
||||
} from "./Prompt";
|
||||
import { getBiggestPrompt } from "./PromptHelper";
|
||||
import { Response } from "./Response";
|
||||
import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext";
|
||||
import { Event } from "./callbacks/CallbackManager";
|
||||
import { LLM } from "./llm/LLM";
|
||||
|
||||
/**
|
||||
* Response modes of the response synthesizer
|
||||
@@ -231,6 +231,7 @@ export class TreeSummarize implements BaseResponseBuilder {
|
||||
throw new Error("Must have at least one text chunk");
|
||||
}
|
||||
|
||||
// Should we send the query here too?
|
||||
const packedTextChunks = this.serviceContext.promptHelper.repack(
|
||||
this.summaryTemplate,
|
||||
textChunks,
|
||||
@@ -241,6 +242,7 @@ export class TreeSummarize implements BaseResponseBuilder {
|
||||
await this.serviceContext.llm.complete(
|
||||
this.summaryTemplate({
|
||||
context: packedTextChunks[0],
|
||||
query,
|
||||
}),
|
||||
parentEvent,
|
||||
)
|
||||
@@ -251,6 +253,7 @@ export class TreeSummarize implements BaseResponseBuilder {
|
||||
this.serviceContext.llm.complete(
|
||||
this.summaryTemplate({
|
||||
context: chunk,
|
||||
query,
|
||||
}),
|
||||
parentEvent,
|
||||
),
|
||||
|
||||
@@ -1,11 +1,15 @@
|
||||
import { Event } from "./callbacks/CallbackManager";
|
||||
import { NodeWithScore } from "./Node";
|
||||
import { ServiceContext } from "./ServiceContext";
|
||||
import { Event } from "./callbacks/CallbackManager";
|
||||
|
||||
/**
|
||||
* Retrievers retrieve the nodes that most closely match our query in similarity.
|
||||
*/
|
||||
export interface BaseRetriever {
|
||||
retrieve(query: string, parentEvent?: Event): Promise<NodeWithScore[]>;
|
||||
retrieve(
|
||||
query: string,
|
||||
parentEvent?: Event,
|
||||
preFilters?: unknown,
|
||||
): Promise<NodeWithScore[]>;
|
||||
getServiceContext(): ServiceContext;
|
||||
}
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import { BaseEmbedding, OpenAIEmbedding } from "./Embedding";
|
||||
import { CallbackManager } from "./callbacks/CallbackManager";
|
||||
import { BaseEmbedding, OpenAIEmbedding } from "./embeddings";
|
||||
import { LLM, OpenAI } from "./llm/LLM";
|
||||
import { NodeParser, SimpleNodeParser } from "./NodeParser";
|
||||
import { PromptHelper } from "./PromptHelper";
|
||||
import { CallbackManager } from "./callbacks/CallbackManager";
|
||||
import { LLM, OpenAI } from "./llm/LLM";
|
||||
|
||||
/**
|
||||
* The ServiceContext is a collection of components that are used in different parts of the application.
|
||||
|
||||
@@ -20,7 +20,8 @@ interface BaseCallbackResponse {
|
||||
event: Event;
|
||||
}
|
||||
|
||||
export interface StreamToken {
|
||||
//Specify StreamToken per mainstream LLM
|
||||
export interface DefaultStreamToken {
|
||||
id: string;
|
||||
object: string;
|
||||
created: number;
|
||||
@@ -29,16 +30,34 @@ export interface StreamToken {
|
||||
index: number;
|
||||
delta: {
|
||||
content?: string | null;
|
||||
role?: "user" | "assistant" | "system" | "function";
|
||||
role?: "user" | "assistant" | "system" | "function" | "tool";
|
||||
};
|
||||
finish_reason: string | null;
|
||||
}[];
|
||||
}
|
||||
|
||||
//OpenAI stream token schema is the default.
|
||||
//Note: Anthropic and Replicate also use similar token schemas.
|
||||
export type OpenAIStreamToken = DefaultStreamToken;
|
||||
export type AnthropicStreamToken = {
|
||||
completion: string;
|
||||
model: string;
|
||||
stop_reason: string | undefined;
|
||||
stop?: boolean | undefined;
|
||||
log_id?: string;
|
||||
};
|
||||
|
||||
//
|
||||
//Callback Responses
|
||||
//
|
||||
//TODO: Write Embedding Callbacks
|
||||
|
||||
//StreamCallbackResponse should let practitioners implement callbacks out of the box...
|
||||
//When custom streaming LLMs are involved, people are expected to write their own StreamCallbackResponses
|
||||
export interface StreamCallbackResponse extends BaseCallbackResponse {
|
||||
index: number;
|
||||
isDone?: boolean;
|
||||
token?: StreamToken;
|
||||
token?: DefaultStreamToken;
|
||||
}
|
||||
|
||||
export interface RetrievalCallbackResponse extends BaseCallbackResponse {
|
||||
|
||||
@@ -1,45 +0,0 @@
|
||||
import { ChatCompletionChunk } from "openai/resources/chat";
|
||||
import { Stream } from "openai/streaming";
|
||||
import { globalsHelper } from "../../GlobalsHelper";
|
||||
import { MessageType } from "../../llm/LLM";
|
||||
import { Event, StreamCallbackResponse } from "../CallbackManager";
|
||||
|
||||
/**
|
||||
* Handles the OpenAI streaming interface and pipes it to the callback function
|
||||
* @param response - The response from the OpenAI API.
|
||||
* @param onLLMStream - A callback function to handle the LLM stream.
|
||||
* @param parentEvent - An optional parent event.
|
||||
* @returns A promise that resolves to an object with a message and a role.
|
||||
*/
|
||||
export async function handleOpenAIStream({
|
||||
response,
|
||||
onLLMStream,
|
||||
parentEvent,
|
||||
}: {
|
||||
response: Stream<ChatCompletionChunk>;
|
||||
onLLMStream: (data: StreamCallbackResponse) => void;
|
||||
parentEvent?: Event;
|
||||
}): Promise<{ message: string; role: MessageType }> {
|
||||
const event = globalsHelper.createEvent({
|
||||
parentEvent,
|
||||
type: "llmPredict",
|
||||
});
|
||||
let index = 0;
|
||||
let cumulativeText = "";
|
||||
let messageRole: MessageType = "assistant";
|
||||
for await (const part of response) {
|
||||
const { content = "", role = "assistant" } = part.choices[0].delta;
|
||||
|
||||
// ignore the first token
|
||||
if (!content && role === "assistant" && index === 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
cumulativeText += content;
|
||||
messageRole = role;
|
||||
onLLMStream?.({ event, index, token: part });
|
||||
index++;
|
||||
}
|
||||
onLLMStream?.({ event, index, isDone: true });
|
||||
return { message: cumulativeText, role: messageRole };
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
import { MultiModalEmbedding } from "./MultiModalEmbedding";
|
||||
import { ImageType, readImage } from "./utils";
|
||||
|
||||
export enum ClipEmbeddingModelType {
|
||||
XENOVA_CLIP_VIT_BASE_PATCH32 = "Xenova/clip-vit-base-patch32",
|
||||
XENOVA_CLIP_VIT_BASE_PATCH16 = "Xenova/clip-vit-base-patch16",
|
||||
}
|
||||
|
||||
export class ClipEmbedding extends MultiModalEmbedding {
|
||||
modelType: ClipEmbeddingModelType =
|
||||
ClipEmbeddingModelType.XENOVA_CLIP_VIT_BASE_PATCH16;
|
||||
|
||||
private tokenizer: any;
|
||||
private processor: any;
|
||||
private visionModel: any;
|
||||
private textModel: any;
|
||||
|
||||
async getTokenizer() {
|
||||
if (!this.tokenizer) {
|
||||
const { AutoTokenizer } = await import("@xenova/transformers");
|
||||
this.tokenizer = await AutoTokenizer.from_pretrained(this.modelType);
|
||||
}
|
||||
return this.tokenizer;
|
||||
}
|
||||
|
||||
async getProcessor() {
|
||||
if (!this.processor) {
|
||||
const { AutoProcessor } = await import("@xenova/transformers");
|
||||
this.processor = await AutoProcessor.from_pretrained(this.modelType);
|
||||
}
|
||||
return this.processor;
|
||||
}
|
||||
|
||||
async getVisionModel() {
|
||||
if (!this.visionModel) {
|
||||
const { CLIPVisionModelWithProjection } = await import(
|
||||
"@xenova/transformers"
|
||||
);
|
||||
this.visionModel = await CLIPVisionModelWithProjection.from_pretrained(
|
||||
this.modelType,
|
||||
);
|
||||
}
|
||||
|
||||
return this.visionModel;
|
||||
}
|
||||
|
||||
async getTextModel() {
|
||||
if (!this.textModel) {
|
||||
const { CLIPTextModelWithProjection } = await import(
|
||||
"@xenova/transformers"
|
||||
);
|
||||
this.textModel = await CLIPTextModelWithProjection.from_pretrained(
|
||||
this.modelType,
|
||||
);
|
||||
}
|
||||
|
||||
return this.textModel;
|
||||
}
|
||||
|
||||
async getImageEmbedding(image: ImageType): Promise<number[]> {
|
||||
const loadedImage = await readImage(image);
|
||||
const imageInputs = await (await this.getProcessor())(loadedImage);
|
||||
const { image_embeds } = await (await this.getVisionModel())(imageInputs);
|
||||
return image_embeds.data;
|
||||
}
|
||||
|
||||
async getTextEmbedding(text: string): Promise<number[]> {
|
||||
const textInputs = await (
|
||||
await this.getTokenizer()
|
||||
)([text], { padding: true, truncation: true });
|
||||
const { text_embeds } = await (await this.getTextModel())(textInputs);
|
||||
return text_embeds.data;
|
||||
}
|
||||
|
||||
async getQueryEmbedding(query: string): Promise<number[]> {
|
||||
return this.getTextEmbedding(query);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
import { BaseEmbedding } from "./types";
|
||||
import { ImageType } from "./utils";
|
||||
|
||||
/*
|
||||
* Base class for Multi Modal embeddings.
|
||||
*/
|
||||
|
||||
export abstract class MultiModalEmbedding extends BaseEmbedding {
|
||||
abstract getImageEmbedding(images: ImageType): Promise<number[]>;
|
||||
|
||||
async getImageEmbeddings(images: ImageType[]): Promise<number[][]> {
|
||||
// Embed the input sequence of images asynchronously.
|
||||
return Promise.all(
|
||||
images.map((imgFilePath) => this.getImageEmbedding(imgFilePath)),
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
import { ClientOptions as OpenAIClientOptions } from "openai";
|
||||
import {
|
||||
AzureOpenAIConfig,
|
||||
getAzureBaseUrl,
|
||||
getAzureConfigFromEnv,
|
||||
getAzureModel,
|
||||
shouldUseAzure,
|
||||
} from "../llm/azure";
|
||||
import { OpenAISession, getOpenAISession } from "../llm/openai";
|
||||
import { BaseEmbedding } from "./types";
|
||||
|
||||
export enum OpenAIEmbeddingModelType {
|
||||
TEXT_EMBED_ADA_002 = "text-embedding-ada-002",
|
||||
}
|
||||
|
||||
export class OpenAIEmbedding extends BaseEmbedding {
|
||||
model: OpenAIEmbeddingModelType;
|
||||
|
||||
// OpenAI session params
|
||||
apiKey?: string = undefined;
|
||||
maxRetries: number;
|
||||
timeout?: number;
|
||||
additionalSessionOptions?: Omit<
|
||||
Partial<OpenAIClientOptions>,
|
||||
"apiKey" | "maxRetries" | "timeout"
|
||||
>;
|
||||
|
||||
session: OpenAISession;
|
||||
|
||||
constructor(init?: Partial<OpenAIEmbedding> & { azure?: AzureOpenAIConfig }) {
|
||||
super();
|
||||
|
||||
this.model = OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002;
|
||||
|
||||
this.maxRetries = init?.maxRetries ?? 10;
|
||||
this.timeout = init?.timeout ?? 60 * 1000; // Default is 60 seconds
|
||||
this.additionalSessionOptions = init?.additionalSessionOptions;
|
||||
|
||||
if (init?.azure || shouldUseAzure()) {
|
||||
const azureConfig = getAzureConfigFromEnv({
|
||||
...init?.azure,
|
||||
model: getAzureModel(this.model),
|
||||
});
|
||||
|
||||
if (!azureConfig.apiKey) {
|
||||
throw new Error(
|
||||
"Azure API key is required for OpenAI Azure models. Please set the AZURE_OPENAI_KEY environment variable.",
|
||||
);
|
||||
}
|
||||
|
||||
this.apiKey = azureConfig.apiKey;
|
||||
this.session =
|
||||
init?.session ??
|
||||
getOpenAISession({
|
||||
azure: true,
|
||||
apiKey: this.apiKey,
|
||||
baseURL: getAzureBaseUrl(azureConfig),
|
||||
maxRetries: this.maxRetries,
|
||||
timeout: this.timeout,
|
||||
defaultQuery: { "api-version": azureConfig.apiVersion },
|
||||
...this.additionalSessionOptions,
|
||||
});
|
||||
} else {
|
||||
this.apiKey = init?.apiKey ?? undefined;
|
||||
this.session =
|
||||
init?.session ??
|
||||
getOpenAISession({
|
||||
apiKey: this.apiKey,
|
||||
maxRetries: this.maxRetries,
|
||||
timeout: this.timeout,
|
||||
...this.additionalSessionOptions,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private async getOpenAIEmbedding(input: string) {
|
||||
const { data } = await this.session.openai.embeddings.create({
|
||||
model: this.model,
|
||||
input,
|
||||
});
|
||||
|
||||
return data[0].embedding;
|
||||
}
|
||||
|
||||
async getTextEmbedding(text: string): Promise<number[]> {
|
||||
return this.getOpenAIEmbedding(text);
|
||||
}
|
||||
|
||||
async getQueryEmbedding(query: string): Promise<number[]> {
|
||||
return this.getOpenAIEmbedding(query);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
export * from "./ClipEmbedding";
|
||||
export * from "./MultiModalEmbedding";
|
||||
export * from "./OpenAIEmbedding";
|
||||
export * from "./types";
|
||||
export * from "./utils";
|
||||
@@ -0,0 +1,24 @@
|
||||
import { similarity } from "./utils";
|
||||
|
||||
/**
|
||||
* Similarity type
|
||||
* Default is cosine similarity. Dot product and negative Euclidean distance are also supported.
|
||||
*/
|
||||
export enum SimilarityType {
|
||||
DEFAULT = "cosine",
|
||||
DOT_PRODUCT = "dot_product",
|
||||
EUCLIDEAN = "euclidean",
|
||||
}
|
||||
|
||||
export abstract class BaseEmbedding {
|
||||
similarity(
|
||||
embedding1: number[],
|
||||
embedding2: number[],
|
||||
mode: SimilarityType = SimilarityType.DEFAULT,
|
||||
): number {
|
||||
return similarity(embedding1, embedding2, mode);
|
||||
}
|
||||
|
||||
abstract getTextEmbedding(text: string): Promise<number[]>;
|
||||
abstract getQueryEmbedding(query: string): Promise<number[]>;
|
||||
}
|
||||
@@ -1,33 +1,16 @@
|
||||
import { ClientOptions as OpenAIClientOptions } from "openai";
|
||||
|
||||
import { DEFAULT_SIMILARITY_TOP_K } from "./constants";
|
||||
import {
|
||||
AzureOpenAIConfig,
|
||||
getAzureBaseUrl,
|
||||
getAzureConfigFromEnv,
|
||||
getAzureModel,
|
||||
shouldUseAzure,
|
||||
} from "./llm/azure";
|
||||
import { OpenAISession, getOpenAISession } from "./llm/openai";
|
||||
import { VectorStoreQueryMode } from "./storage/vectorStore/types";
|
||||
|
||||
/**
|
||||
* Similarity type
|
||||
* Default is cosine similarity. Dot product and negative Euclidean distance are also supported.
|
||||
*/
|
||||
export enum SimilarityType {
|
||||
DEFAULT = "cosine",
|
||||
DOT_PRODUCT = "dot_product",
|
||||
EUCLIDEAN = "euclidean",
|
||||
}
|
||||
import _ from "lodash";
|
||||
import { DEFAULT_SIMILARITY_TOP_K } from "../constants";
|
||||
import { VectorStoreQueryMode } from "../storage";
|
||||
import { SimilarityType } from "./types";
|
||||
|
||||
/**
|
||||
* The similarity between two embeddings.
|
||||
* @param embedding1
|
||||
* @param embedding2
|
||||
* @param mode
|
||||
* @returns similartiy score with higher numbers meaning the two embeddings are more similar
|
||||
* @returns similarity score with higher numbers meaning the two embeddings are more similar
|
||||
*/
|
||||
|
||||
export function similarity(
|
||||
embedding1: number[],
|
||||
embedding2: number[],
|
||||
@@ -42,7 +25,6 @@ export function similarity(
|
||||
// will probably cause some avoidable loss of floating point precision
|
||||
// ml-distance is worth watching although they currently also use the naive
|
||||
// formulas
|
||||
|
||||
function norm(x: number[]): number {
|
||||
let result = 0;
|
||||
for (let i = 0; i < x.length; i++) {
|
||||
@@ -201,98 +183,14 @@ export function getTopKMMREmbeddings(
|
||||
|
||||
return [resultSimilarities, resultIds];
|
||||
}
|
||||
|
||||
export abstract class BaseEmbedding {
|
||||
similarity(
|
||||
embedding1: number[],
|
||||
embedding2: number[],
|
||||
mode: SimilarityType = SimilarityType.DEFAULT,
|
||||
): number {
|
||||
return similarity(embedding1, embedding2, mode);
|
||||
}
|
||||
|
||||
abstract getTextEmbedding(text: string): Promise<number[]>;
|
||||
abstract getQueryEmbedding(query: string): Promise<number[]>;
|
||||
}
|
||||
|
||||
enum OpenAIEmbeddingModelType {
|
||||
TEXT_EMBED_ADA_002 = "text-embedding-ada-002",
|
||||
}
|
||||
|
||||
export class OpenAIEmbedding extends BaseEmbedding {
|
||||
model: OpenAIEmbeddingModelType;
|
||||
|
||||
// OpenAI session params
|
||||
apiKey?: string = undefined;
|
||||
maxRetries: number;
|
||||
timeout?: number;
|
||||
additionalSessionOptions?: Omit<
|
||||
Partial<OpenAIClientOptions>,
|
||||
"apiKey" | "maxRetries" | "timeout"
|
||||
>;
|
||||
|
||||
session: OpenAISession;
|
||||
|
||||
constructor(init?: Partial<OpenAIEmbedding> & { azure?: AzureOpenAIConfig }) {
|
||||
super();
|
||||
|
||||
this.model = OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002;
|
||||
|
||||
this.maxRetries = init?.maxRetries ?? 10;
|
||||
this.timeout = init?.timeout ?? 60 * 1000; // Default is 60 seconds
|
||||
this.additionalSessionOptions = init?.additionalSessionOptions;
|
||||
|
||||
if (init?.azure || shouldUseAzure()) {
|
||||
const azureConfig = getAzureConfigFromEnv({
|
||||
...init?.azure,
|
||||
model: getAzureModel(this.model),
|
||||
});
|
||||
|
||||
if (!azureConfig.apiKey) {
|
||||
throw new Error(
|
||||
"Azure API key is required for OpenAI Azure models. Please set the AZURE_OPENAI_KEY environment variable.",
|
||||
);
|
||||
}
|
||||
|
||||
this.apiKey = azureConfig.apiKey;
|
||||
this.session =
|
||||
init?.session ??
|
||||
getOpenAISession({
|
||||
azure: true,
|
||||
apiKey: this.apiKey,
|
||||
baseURL: getAzureBaseUrl(azureConfig),
|
||||
maxRetries: this.maxRetries,
|
||||
timeout: this.timeout,
|
||||
defaultQuery: { "api-version": azureConfig.apiVersion },
|
||||
...this.additionalSessionOptions,
|
||||
});
|
||||
} else {
|
||||
this.apiKey = init?.apiKey ?? undefined;
|
||||
this.session =
|
||||
init?.session ??
|
||||
getOpenAISession({
|
||||
apiKey: this.apiKey,
|
||||
maxRetries: this.maxRetries,
|
||||
timeout: this.timeout,
|
||||
...this.additionalSessionOptions,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private async getOpenAIEmbedding(input: string) {
|
||||
const { data } = await this.session.openai.embeddings.create({
|
||||
model: this.model,
|
||||
input,
|
||||
});
|
||||
|
||||
return data[0].embedding;
|
||||
}
|
||||
|
||||
async getTextEmbedding(text: string): Promise<number[]> {
|
||||
return this.getOpenAIEmbedding(text);
|
||||
}
|
||||
|
||||
async getQueryEmbedding(query: string): Promise<number[]> {
|
||||
return this.getOpenAIEmbedding(query);
|
||||
export async function readImage(input: ImageType) {
|
||||
const { RawImage } = await import("@xenova/transformers");
|
||||
if (input instanceof Blob) {
|
||||
return await RawImage.fromBlob(input);
|
||||
} else if (_.isString(input) || input instanceof URL) {
|
||||
return await RawImage.fromURL(input);
|
||||
} else {
|
||||
throw new Error(`Unsupported input type: ${typeof input}`);
|
||||
}
|
||||
}
|
||||
export type ImageType = string | Blob | URL;
|
||||
@@ -1,10 +1,11 @@
|
||||
export * from "./ChatEngine";
|
||||
export * from "./Embedding";
|
||||
export * from "./ChatHistory";
|
||||
export * from "./GlobalsHelper";
|
||||
export * from "./Node";
|
||||
export * from "./NodeParser";
|
||||
export * from "./OutputParser";
|
||||
export * from "./Prompt";
|
||||
export * from "./PromptHelper";
|
||||
export * from "./QueryEngine";
|
||||
export * from "./QuestionGenerator";
|
||||
export * from "./Response";
|
||||
@@ -13,18 +14,17 @@ export * from "./Retriever";
|
||||
export * from "./ServiceContext";
|
||||
export * from "./TextSplitter";
|
||||
export * from "./Tool";
|
||||
export * from "./constants";
|
||||
export * from "./llm/LLM";
|
||||
|
||||
export * from "./indices";
|
||||
|
||||
export * from "./callbacks/CallbackManager";
|
||||
|
||||
export * from "./constants";
|
||||
export * from "./embeddings";
|
||||
export * from "./indices";
|
||||
export * from "./llm/LLM";
|
||||
export * from "./readers/CSVReader";
|
||||
export * from "./readers/HTMLReader";
|
||||
export * from "./readers/MarkdownReader";
|
||||
export * from "./readers/NotionReader";
|
||||
export * from "./readers/PDFReader";
|
||||
export * from "./readers/SimpleDirectoryReader";
|
||||
export * from "./readers/SimpleMongoReader";
|
||||
export * from "./readers/base";
|
||||
|
||||
export * from "./storage";
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
import { NodeWithScore } from "../Node";
|
||||
|
||||
export interface BaseNodePostprocessor {
|
||||
postprocessNodes: (nodes: NodeWithScore[]) => NodeWithScore[];
|
||||
}
|
||||
|
||||
export class SimilarityPostprocessor implements BaseNodePostprocessor {
|
||||
similarityCutoff?: number;
|
||||
|
||||
constructor(options?: { similarityCutoff?: number }) {
|
||||
this.similarityCutoff = options?.similarityCutoff;
|
||||
}
|
||||
|
||||
postprocessNodes(nodes: NodeWithScore[]) {
|
||||
if (this.similarityCutoff === undefined) return nodes;
|
||||
|
||||
const cutoff = this.similarityCutoff || 0;
|
||||
return nodes.filter((node) => node.score && node.score >= cutoff);
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
export * from "./BaseIndex";
|
||||
export * from "./BaseNodePostprocessor";
|
||||
export * from "./keyword";
|
||||
export * from "./summary";
|
||||
export * from "./vectorStore";
|
||||
|
||||
@@ -15,6 +15,7 @@ import {
|
||||
IndexStructType,
|
||||
KeywordTable,
|
||||
} from "../BaseIndex";
|
||||
import { BaseNodePostprocessor } from "../BaseNodePostprocessor";
|
||||
import {
|
||||
KeywordTableLLMRetriever,
|
||||
KeywordTableRAKERetriever,
|
||||
@@ -129,11 +130,15 @@ export class KeywordTableIndex extends BaseIndex<KeywordTable> {
|
||||
asQueryEngine(options?: {
|
||||
retriever?: BaseRetriever;
|
||||
responseSynthesizer?: ResponseSynthesizer;
|
||||
preFilters?: unknown;
|
||||
nodePostprocessors?: BaseNodePostprocessor[];
|
||||
}): BaseQueryEngine {
|
||||
const { retriever, responseSynthesizer } = options ?? {};
|
||||
return new RetrieverQueryEngine(
|
||||
retriever ?? this.asRetriever(),
|
||||
responseSynthesizer,
|
||||
options?.preFilters,
|
||||
options?.nodePostprocessors,
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -10,17 +10,18 @@ import {
|
||||
ServiceContext,
|
||||
serviceContextFromDefaults,
|
||||
} from "../../ServiceContext";
|
||||
import { BaseDocumentStore, RefDocInfo } from "../../storage/docStore/types";
|
||||
import {
|
||||
StorageContext,
|
||||
storageContextFromDefaults,
|
||||
} from "../../storage/StorageContext";
|
||||
import { BaseDocumentStore, RefDocInfo } from "../../storage/docStore/types";
|
||||
import {
|
||||
BaseIndex,
|
||||
BaseIndexInit,
|
||||
IndexList,
|
||||
IndexStructType,
|
||||
} from "../BaseIndex";
|
||||
import { BaseNodePostprocessor } from "../BaseNodePostprocessor";
|
||||
import {
|
||||
SummaryIndexLLMRetriever,
|
||||
SummaryIndexRetriever,
|
||||
@@ -155,6 +156,8 @@ export class SummaryIndex extends BaseIndex<IndexList> {
|
||||
asQueryEngine(options?: {
|
||||
retriever?: BaseRetriever;
|
||||
responseSynthesizer?: ResponseSynthesizer;
|
||||
preFilters?: unknown;
|
||||
nodePostprocessors?: BaseNodePostprocessor[];
|
||||
}): BaseQueryEngine {
|
||||
let { retriever, responseSynthesizer } = options ?? {};
|
||||
|
||||
@@ -170,7 +173,12 @@ export class SummaryIndex extends BaseIndex<IndexList> {
|
||||
});
|
||||
}
|
||||
|
||||
return new RetrieverQueryEngine(retriever, responseSynthesizer);
|
||||
return new RetrieverQueryEngine(
|
||||
retriever,
|
||||
responseSynthesizer,
|
||||
options?.preFilters,
|
||||
options?.nodePostprocessors,
|
||||
);
|
||||
}
|
||||
|
||||
static async buildIndexFromNodes(
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
import { Event } from "../../callbacks/CallbackManager";
|
||||
import { DEFAULT_SIMILARITY_TOP_K } from "../../constants";
|
||||
import { globalsHelper } from "../../GlobalsHelper";
|
||||
import { NodeWithScore } from "../../Node";
|
||||
import { BaseRetriever } from "../../Retriever";
|
||||
import { ServiceContext } from "../../ServiceContext";
|
||||
import { Event } from "../../callbacks/CallbackManager";
|
||||
import { DEFAULT_SIMILARITY_TOP_K } from "../../constants";
|
||||
import {
|
||||
VectorStoreQuery,
|
||||
VectorStoreQueryMode,
|
||||
@@ -32,7 +32,11 @@ export class VectorIndexRetriever implements BaseRetriever {
|
||||
this.similarityTopK = similarityTopK ?? DEFAULT_SIMILARITY_TOP_K;
|
||||
}
|
||||
|
||||
async retrieve(query: string, parentEvent?: Event): Promise<NodeWithScore[]> {
|
||||
async retrieve(
|
||||
query: string,
|
||||
parentEvent?: Event,
|
||||
preFilters?: unknown,
|
||||
): Promise<NodeWithScore[]> {
|
||||
const queryEmbedding =
|
||||
await this.serviceContext.embedModel.getQueryEmbedding(query);
|
||||
|
||||
@@ -41,10 +45,15 @@ export class VectorIndexRetriever implements BaseRetriever {
|
||||
mode: VectorStoreQueryMode.DEFAULT,
|
||||
similarityTopK: this.similarityTopK,
|
||||
};
|
||||
const result = await this.index.vectorStore.query(q);
|
||||
const result = await this.index.vectorStore.query(q, preFilters);
|
||||
|
||||
let nodesWithScores: NodeWithScore[] = [];
|
||||
for (let i = 0; i < result.ids.length; i++) {
|
||||
const nodeFromResult = result.nodes?.[i];
|
||||
if (!this.index.indexStruct.nodesDict[result.ids[i]] && nodeFromResult) {
|
||||
this.index.indexStruct.nodesDict[result.ids[i]] = nodeFromResult;
|
||||
}
|
||||
|
||||
const node = this.index.indexStruct.nodesDict[result.ids[i]];
|
||||
nodesWithScores.push({
|
||||
node: node,
|
||||
|
||||
@@ -18,6 +18,7 @@ import {
|
||||
IndexDict,
|
||||
IndexStructType,
|
||||
} from "../BaseIndex";
|
||||
import { BaseNodePostprocessor } from "../BaseNodePostprocessor";
|
||||
import { VectorIndexRetriever } from "./VectorIndexRetriever";
|
||||
|
||||
export interface VectorIndexOptions {
|
||||
@@ -87,24 +88,23 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
|
||||
);
|
||||
}
|
||||
|
||||
if (!indexStruct && !options.nodes) {
|
||||
if (options.nodes) {
|
||||
// If nodes are passed in, then we need to update the index
|
||||
indexStruct = await VectorStoreIndex.buildIndexFromNodes(
|
||||
options.nodes,
|
||||
serviceContext,
|
||||
vectorStore,
|
||||
docStore,
|
||||
indexStruct,
|
||||
);
|
||||
|
||||
await indexStore.addIndexStruct(indexStruct);
|
||||
} else if (!indexStruct) {
|
||||
throw new Error(
|
||||
"Cannot initialize VectorStoreIndex without nodes or indexStruct",
|
||||
);
|
||||
}
|
||||
|
||||
const nodes = options.nodes ?? [];
|
||||
|
||||
indexStruct = await VectorStoreIndex.buildIndexFromNodes(
|
||||
nodes,
|
||||
serviceContext,
|
||||
vectorStore,
|
||||
docStore,
|
||||
indexStruct,
|
||||
);
|
||||
|
||||
await indexStore.addIndexStruct(indexStruct);
|
||||
|
||||
return new VectorStoreIndex({
|
||||
storageContext,
|
||||
serviceContext,
|
||||
@@ -219,6 +219,27 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
|
||||
return index;
|
||||
}
|
||||
|
||||
static async fromVectorStore(
|
||||
vectorStore: VectorStore,
|
||||
serviceContext: ServiceContext,
|
||||
) {
|
||||
if (!vectorStore.storesText) {
|
||||
throw new Error(
|
||||
"Cannot initialize from a vector store that does not store text",
|
||||
);
|
||||
}
|
||||
|
||||
const storageContext = await storageContextFromDefaults({ vectorStore });
|
||||
|
||||
const index = await VectorStoreIndex.init({
|
||||
nodes: [],
|
||||
storageContext,
|
||||
serviceContext,
|
||||
});
|
||||
|
||||
return index;
|
||||
}
|
||||
|
||||
asRetriever(options?: any): VectorIndexRetriever {
|
||||
return new VectorIndexRetriever({ index: this, ...options });
|
||||
}
|
||||
@@ -226,11 +247,15 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
|
||||
asQueryEngine(options?: {
|
||||
retriever?: BaseRetriever;
|
||||
responseSynthesizer?: ResponseSynthesizer;
|
||||
preFilters?: unknown;
|
||||
nodePostprocessors?: BaseNodePostprocessor[];
|
||||
}): BaseQueryEngine {
|
||||
const { retriever, responseSynthesizer } = options ?? {};
|
||||
return new RetrieverQueryEngine(
|
||||
retriever ?? this.asRetriever(),
|
||||
responseSynthesizer,
|
||||
options?.preFilters,
|
||||
options?.nodePostprocessors,
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
+474
-73
@@ -1,6 +1,16 @@
|
||||
import OpenAILLM, { ClientOptions as OpenAIClientOptions } from "openai";
|
||||
import { CallbackManager, Event } from "../callbacks/CallbackManager";
|
||||
import { handleOpenAIStream } from "../callbacks/utility/handleOpenAIStream";
|
||||
import {
|
||||
AnthropicStreamToken,
|
||||
CallbackManager,
|
||||
Event,
|
||||
EventType,
|
||||
OpenAIStreamToken,
|
||||
StreamCallbackResponse,
|
||||
} from "../callbacks/CallbackManager";
|
||||
|
||||
import { ChatCompletionMessageParam } from "openai/resources";
|
||||
import { LLMOptions } from "portkey-ai";
|
||||
import { globalsHelper, Tokenizers } from "../GlobalsHelper";
|
||||
import {
|
||||
ANTHROPIC_AI_PROMPT,
|
||||
ANTHROPIC_HUMAN_PROMPT,
|
||||
@@ -14,7 +24,8 @@ import {
|
||||
getAzureModel,
|
||||
shouldUseAzure,
|
||||
} from "./azure";
|
||||
import { OpenAISession, getOpenAISession } from "./openai";
|
||||
import { getOpenAISession, OpenAISession } from "./openai";
|
||||
import { getPortkeySession, PortkeySession } from "./portkey";
|
||||
import { ReplicateSession } from "./replicate";
|
||||
|
||||
export type MessageType =
|
||||
@@ -22,10 +33,11 @@ export type MessageType =
|
||||
| "assistant"
|
||||
| "system"
|
||||
| "generic"
|
||||
| "function";
|
||||
| "function"
|
||||
| "memory";
|
||||
|
||||
export interface ChatMessage {
|
||||
content: string;
|
||||
content: any;
|
||||
role: MessageType;
|
||||
}
|
||||
|
||||
@@ -38,31 +50,67 @@ export interface ChatResponse {
|
||||
// NOTE in case we need CompletionResponse to diverge from ChatResponse in the future
|
||||
export type CompletionResponse = ChatResponse;
|
||||
|
||||
export interface LLMMetadata {
|
||||
model: string;
|
||||
temperature: number;
|
||||
topP: number;
|
||||
maxTokens?: number;
|
||||
contextWindow: number;
|
||||
tokenizer: Tokenizers | undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Unified language model interface
|
||||
*/
|
||||
export interface LLM {
|
||||
metadata: LLMMetadata;
|
||||
// Whether a LLM has streaming support
|
||||
hasStreaming: boolean;
|
||||
/**
|
||||
* Get a chat response from the LLM
|
||||
* @param messages
|
||||
*
|
||||
* The return type of chat() and complete() are set by the "streaming" parameter being set to True.
|
||||
*/
|
||||
chat(messages: ChatMessage[], parentEvent?: Event): Promise<ChatResponse>;
|
||||
chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(
|
||||
messages: ChatMessage[],
|
||||
parentEvent?: Event,
|
||||
streaming?: T,
|
||||
): Promise<R>;
|
||||
|
||||
/**
|
||||
* Get a prompt completion from the LLM
|
||||
* @param prompt the prompt to complete
|
||||
*/
|
||||
complete(prompt: string, parentEvent?: Event): Promise<CompletionResponse>;
|
||||
complete<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(
|
||||
prompt: string,
|
||||
parentEvent?: Event,
|
||||
streaming?: T,
|
||||
): Promise<R>;
|
||||
|
||||
/**
|
||||
* Calculates the number of tokens needed for the given chat messages
|
||||
*/
|
||||
tokens(messages: ChatMessage[]): number;
|
||||
}
|
||||
|
||||
export const GPT4_MODELS = {
|
||||
"gpt-4": { contextWindow: 8192 },
|
||||
"gpt-4-32k": { contextWindow: 32768 },
|
||||
"gpt-4-1106-preview": { contextWindow: 128000 },
|
||||
"gpt-4-vision-preview": { contextWindow: 8192 },
|
||||
};
|
||||
|
||||
export const TURBO_MODELS = {
|
||||
export const GPT35_MODELS = {
|
||||
"gpt-3.5-turbo": { contextWindow: 4096 },
|
||||
"gpt-3.5-turbo-16k": { contextWindow: 16384 },
|
||||
"gpt-3.5-turbo-1106": { contextWindow: 16384 },
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -70,20 +118,22 @@ export const TURBO_MODELS = {
|
||||
*/
|
||||
export const ALL_AVAILABLE_OPENAI_MODELS = {
|
||||
...GPT4_MODELS,
|
||||
...TURBO_MODELS,
|
||||
...GPT35_MODELS,
|
||||
};
|
||||
|
||||
/**
|
||||
* OpenAI LLM implementation
|
||||
*/
|
||||
export class OpenAI implements LLM {
|
||||
hasStreaming: boolean = true;
|
||||
|
||||
// Per completion OpenAI params
|
||||
model: keyof typeof ALL_AVAILABLE_OPENAI_MODELS;
|
||||
temperature: number;
|
||||
topP: number;
|
||||
maxTokens?: number;
|
||||
additionalChatOptions?: Omit<
|
||||
Partial<OpenAILLM.Chat.CompletionCreateParams>,
|
||||
Partial<OpenAILLM.Chat.ChatCompletionCreateParams>,
|
||||
"max_tokens" | "messages" | "model" | "temperature" | "top_p" | "streaming"
|
||||
>;
|
||||
|
||||
@@ -153,6 +203,32 @@ export class OpenAI implements LLM {
|
||||
this.callbackManager = init?.callbackManager;
|
||||
}
|
||||
|
||||
get metadata() {
|
||||
return {
|
||||
model: this.model,
|
||||
temperature: this.temperature,
|
||||
topP: this.topP,
|
||||
maxTokens: this.maxTokens,
|
||||
contextWindow: ALL_AVAILABLE_OPENAI_MODELS[this.model].contextWindow,
|
||||
tokenizer: Tokenizers.CL100K_BASE,
|
||||
};
|
||||
}
|
||||
|
||||
tokens(messages: ChatMessage[]): number {
|
||||
// for latest OpenAI models, see https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
|
||||
const tokenizer = globalsHelper.tokenizer(this.metadata.tokenizer);
|
||||
const tokensPerMessage = 3;
|
||||
let numTokens = 0;
|
||||
for (const message of messages) {
|
||||
numTokens += tokensPerMessage;
|
||||
for (const value of Object.values(message)) {
|
||||
numTokens += tokenizer(value).length;
|
||||
}
|
||||
}
|
||||
numTokens += 3; // every reply is primed with <|im_start|>assistant<|im_sep|>
|
||||
return numTokens;
|
||||
}
|
||||
|
||||
mapMessageType(
|
||||
messageType: MessageType,
|
||||
): "user" | "assistant" | "system" | "function" {
|
||||
@@ -170,52 +246,124 @@ export class OpenAI implements LLM {
|
||||
}
|
||||
}
|
||||
|
||||
async chat(
|
||||
messages: ChatMessage[],
|
||||
parentEvent?: Event,
|
||||
): Promise<ChatResponse> {
|
||||
const baseRequestParams: OpenAILLM.Chat.CompletionCreateParams = {
|
||||
async chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(messages: ChatMessage[], parentEvent?: Event, streaming?: T): Promise<R> {
|
||||
const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
|
||||
model: this.model,
|
||||
temperature: this.temperature,
|
||||
max_tokens: this.maxTokens,
|
||||
messages: messages.map((message) => ({
|
||||
role: this.mapMessageType(message.role),
|
||||
content: message.content,
|
||||
})),
|
||||
messages: messages.map(
|
||||
(message) =>
|
||||
({
|
||||
role: this.mapMessageType(message.role),
|
||||
content: message.content,
|
||||
}) as ChatCompletionMessageParam,
|
||||
),
|
||||
top_p: this.topP,
|
||||
...this.additionalChatOptions,
|
||||
};
|
||||
// Streaming
|
||||
if (streaming) {
|
||||
if (!this.hasStreaming) {
|
||||
throw Error("No streaming support for this LLM.");
|
||||
}
|
||||
return this.streamChat(messages, parentEvent) as R;
|
||||
}
|
||||
// Non-streaming
|
||||
const response = await this.session.openai.chat.completions.create({
|
||||
...baseRequestParams,
|
||||
stream: false,
|
||||
});
|
||||
|
||||
const content = response.choices[0].message?.content ?? "";
|
||||
return {
|
||||
message: { content, role: response.choices[0].message.role },
|
||||
} as R;
|
||||
}
|
||||
|
||||
async complete<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(prompt: string, parentEvent?: Event, streaming?: T): Promise<R> {
|
||||
return this.chat(
|
||||
[{ content: prompt, role: "user" }],
|
||||
parentEvent,
|
||||
streaming,
|
||||
);
|
||||
}
|
||||
|
||||
//We can wrap a stream in a generator to add some additional logging behavior
|
||||
//For future edits: syntax for generator type is <typeof Yield, typeof Return, typeof Accept>
|
||||
//"typeof Accept" refers to what types you'll accept when you manually call generator.next(<AcceptType>)
|
||||
protected async *streamChat(
|
||||
messages: ChatMessage[],
|
||||
parentEvent?: Event,
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
|
||||
model: this.model,
|
||||
temperature: this.temperature,
|
||||
max_tokens: this.maxTokens,
|
||||
messages: messages.map(
|
||||
(message) =>
|
||||
({
|
||||
role: this.mapMessageType(message.role),
|
||||
content: message.content,
|
||||
}) as ChatCompletionMessageParam,
|
||||
),
|
||||
top_p: this.topP,
|
||||
...this.additionalChatOptions,
|
||||
};
|
||||
|
||||
if (this.callbackManager?.onLLMStream) {
|
||||
// Streaming
|
||||
const response = await this.session.openai.chat.completions.create({
|
||||
//Now let's wrap our stream in a callback
|
||||
const onLLMStream = this.callbackManager?.onLLMStream
|
||||
? this.callbackManager.onLLMStream
|
||||
: () => {};
|
||||
|
||||
const chunk_stream: AsyncIterable<OpenAIStreamToken> =
|
||||
await this.session.openai.chat.completions.create({
|
||||
...baseRequestParams,
|
||||
stream: true,
|
||||
});
|
||||
|
||||
const { message, role } = await handleOpenAIStream({
|
||||
response,
|
||||
onLLMStream: this.callbackManager.onLLMStream,
|
||||
parentEvent,
|
||||
});
|
||||
return { message: { content: message, role } };
|
||||
} else {
|
||||
// Non-streaming
|
||||
const response = await this.session.openai.chat.completions.create({
|
||||
...baseRequestParams,
|
||||
stream: false,
|
||||
});
|
||||
const event: Event = parentEvent
|
||||
? parentEvent
|
||||
: {
|
||||
id: "unspecified",
|
||||
type: "llmPredict" as EventType,
|
||||
};
|
||||
|
||||
const content = response.choices[0].message?.content ?? "";
|
||||
return { message: { content, role: response.choices[0].message.role } };
|
||||
//Indices
|
||||
var idx_counter: number = 0;
|
||||
for await (const part of chunk_stream) {
|
||||
//Increment
|
||||
part.choices[0].index = idx_counter;
|
||||
const is_done: boolean =
|
||||
part.choices[0].finish_reason === "stop" ? true : false;
|
||||
//onLLMStream Callback
|
||||
|
||||
const stream_callback: StreamCallbackResponse = {
|
||||
event: event,
|
||||
index: idx_counter,
|
||||
isDone: is_done,
|
||||
token: part,
|
||||
};
|
||||
onLLMStream(stream_callback);
|
||||
|
||||
idx_counter++;
|
||||
|
||||
yield part.choices[0].delta.content ? part.choices[0].delta.content : "";
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
async complete(
|
||||
prompt: string,
|
||||
//streamComplete doesn't need to be async because it's child function is already async
|
||||
protected streamComplete(
|
||||
query: string,
|
||||
parentEvent?: Event,
|
||||
): Promise<CompletionResponse> {
|
||||
return this.chat([{ content: prompt, role: "user" }], parentEvent);
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
return this.streamChat([{ content: query, role: "user" }], parentEvent);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -229,10 +377,10 @@ export const ALL_AVAILABLE_LLAMADEUCE_MODELS = {
|
||||
"Llama-2-70b-chat-4bit": {
|
||||
contextWindow: 4096,
|
||||
replicateApi:
|
||||
"replicate/llama70b-v2-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1",
|
||||
"meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3",
|
||||
//^ Model is based off of exllama 4bit.
|
||||
},
|
||||
"Llama-2-13b-chat": {
|
||||
"Llama-2-13b-chat-old": {
|
||||
contextWindow: 4096,
|
||||
replicateApi:
|
||||
"a16z-infra/llama13b-v2-chat:df7690f1994d94e96ad9d568eac121aecf50684a0b0963b25a41cc40061269e5",
|
||||
@@ -241,9 +389,9 @@ export const ALL_AVAILABLE_LLAMADEUCE_MODELS = {
|
||||
"Llama-2-13b-chat-4bit": {
|
||||
contextWindow: 4096,
|
||||
replicateApi:
|
||||
"a16z-infra/llama13b-v2-chat:2a7f981751ec7fdf87b5b91ad4db53683a98082e9ff7bfd12c8cd5ea85980a52",
|
||||
"meta/llama-2-13b-chat:f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
|
||||
},
|
||||
"Llama-2-7b-chat": {
|
||||
"Llama-2-7b-chat-old": {
|
||||
contextWindow: 4096,
|
||||
replicateApi:
|
||||
"a16z-infra/llama7b-v2-chat:4f0a4744c7295c024a1de15e1a63c880d3da035fa1f49bfd344fe076074c8eea",
|
||||
@@ -255,7 +403,7 @@ export const ALL_AVAILABLE_LLAMADEUCE_MODELS = {
|
||||
"Llama-2-7b-chat-4bit": {
|
||||
contextWindow: 4096,
|
||||
replicateApi:
|
||||
"a16z-infra/llama7b-v2-chat:4f0b260b6a13eb53a6b1891f089d57c08f41003ae79458be5011303d81a394dc",
|
||||
"meta/llama-2-7b-chat:13c3cdee13ee059ab779f0291d29054dab00a47dad8261375654de5540165fb0",
|
||||
},
|
||||
};
|
||||
|
||||
@@ -267,6 +415,8 @@ export enum DeuceChatStrategy {
|
||||
// Unfortunately any string only API won't support these properly.
|
||||
REPLICATE4BIT = "replicate4bit",
|
||||
//^ To satisfy Replicate's 4 bit models' requirements where they also insert some INST tags
|
||||
REPLICATE4BITWNEWLINES = "replicate4bitwnewlines",
|
||||
//^ Replicate's documentation recommends using newlines: https://replicate.com/blog/how-to-prompt-llama
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -279,13 +429,14 @@ export class LlamaDeuce implements LLM {
|
||||
topP: number;
|
||||
maxTokens?: number;
|
||||
replicateSession: ReplicateSession;
|
||||
hasStreaming: boolean;
|
||||
|
||||
constructor(init?: Partial<LlamaDeuce>) {
|
||||
this.model = init?.model ?? "Llama-2-70b-chat-4bit";
|
||||
this.chatStrategy =
|
||||
init?.chatStrategy ??
|
||||
(this.model.endsWith("4bit")
|
||||
? DeuceChatStrategy.REPLICATE4BIT // With the newer A16Z/Replicate models they do the system message themselves.
|
||||
? DeuceChatStrategy.REPLICATE4BITWNEWLINES // With the newer Replicate models they do the system message themselves.
|
||||
: DeuceChatStrategy.METAWBOS); // With BOS and EOS seems to work best, although they all have problems past a certain point
|
||||
this.temperature = init?.temperature ?? 0.1; // minimum temperature is 0.01 for Replicate endpoint
|
||||
this.topP = init?.topP ?? 1;
|
||||
@@ -293,6 +444,22 @@ export class LlamaDeuce implements LLM {
|
||||
init?.maxTokens ??
|
||||
ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model].contextWindow; // For Replicate, the default is 500 tokens which is too low.
|
||||
this.replicateSession = init?.replicateSession ?? new ReplicateSession();
|
||||
this.hasStreaming = init?.hasStreaming ?? false;
|
||||
}
|
||||
|
||||
tokens(messages: ChatMessage[]): number {
|
||||
throw new Error("Method not implemented.");
|
||||
}
|
||||
|
||||
get metadata() {
|
||||
return {
|
||||
model: this.model,
|
||||
temperature: this.temperature,
|
||||
topP: this.topP,
|
||||
maxTokens: this.maxTokens,
|
||||
contextWindow: ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model].contextWindow,
|
||||
tokenizer: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
mapMessagesToPrompt(messages: ChatMessage[]) {
|
||||
@@ -303,7 +470,15 @@ export class LlamaDeuce implements LLM {
|
||||
} else if (this.chatStrategy === DeuceChatStrategy.METAWBOS) {
|
||||
return this.mapMessagesToPromptMeta(messages, { withBos: true });
|
||||
} else if (this.chatStrategy === DeuceChatStrategy.REPLICATE4BIT) {
|
||||
return this.mapMessagesToPromptMeta(messages, { replicate4Bit: true });
|
||||
return this.mapMessagesToPromptMeta(messages, {
|
||||
replicate4Bit: true,
|
||||
withNewlines: true,
|
||||
});
|
||||
} else if (this.chatStrategy === DeuceChatStrategy.REPLICATE4BITWNEWLINES) {
|
||||
return this.mapMessagesToPromptMeta(messages, {
|
||||
replicate4Bit: true,
|
||||
withNewlines: true,
|
||||
});
|
||||
} else {
|
||||
return this.mapMessagesToPromptMeta(messages);
|
||||
}
|
||||
@@ -338,9 +513,17 @@ export class LlamaDeuce implements LLM {
|
||||
|
||||
mapMessagesToPromptMeta(
|
||||
messages: ChatMessage[],
|
||||
opts?: { withBos?: boolean; replicate4Bit?: boolean },
|
||||
opts?: {
|
||||
withBos?: boolean;
|
||||
replicate4Bit?: boolean;
|
||||
withNewlines?: boolean;
|
||||
},
|
||||
) {
|
||||
const { withBos = false, replicate4Bit = false } = opts ?? {};
|
||||
const {
|
||||
withBos = false,
|
||||
replicate4Bit = false,
|
||||
withNewlines = false,
|
||||
} = opts ?? {};
|
||||
const DEFAULT_SYSTEM_PROMPT = `You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
|
||||
|
||||
If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.`;
|
||||
@@ -388,21 +571,28 @@ If a question does not make any sense, or is not factually coherent, explain why
|
||||
return {
|
||||
prompt: messages.reduce((acc, message, index) => {
|
||||
if (index % 2 === 0) {
|
||||
return `${acc}${
|
||||
withBos ? BOS : ""
|
||||
}${B_INST} ${message.content.trim()} ${E_INST}`;
|
||||
return (
|
||||
`${acc}${
|
||||
withBos ? BOS : ""
|
||||
}${B_INST} ${message.content.trim()} ${E_INST}` +
|
||||
(withNewlines ? "\n" : "")
|
||||
);
|
||||
} else {
|
||||
return `${acc} ${message.content.trim()} ` + (withBos ? EOS : ""); // Yes, the EOS comes after the space. This is not a mistake.
|
||||
return (
|
||||
`${acc} ${message.content.trim()}` +
|
||||
(withNewlines ? "\n" : " ") +
|
||||
(withBos ? EOS : "")
|
||||
); // Yes, the EOS comes after the space. This is not a mistake.
|
||||
}
|
||||
}, ""),
|
||||
systemPrompt,
|
||||
};
|
||||
}
|
||||
|
||||
async chat(
|
||||
messages: ChatMessage[],
|
||||
_parentEvent?: Event,
|
||||
): Promise<ChatResponse> {
|
||||
async chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(messages: ChatMessage[], _parentEvent?: Event, streaming?: T): Promise<R> {
|
||||
const api = ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model]
|
||||
.replicateApi as `${string}/${string}:${string}`;
|
||||
|
||||
@@ -423,6 +613,9 @@ If a question does not make any sense, or is not factually coherent, explain why
|
||||
replicateOptions.input.max_length = this.maxTokens;
|
||||
}
|
||||
|
||||
//TODO: Add streaming for this
|
||||
|
||||
//Non-streaming
|
||||
const response = await this.replicateSession.replicate.run(
|
||||
api,
|
||||
replicateOptions,
|
||||
@@ -433,24 +626,32 @@ If a question does not make any sense, or is not factually coherent, explain why
|
||||
//^ We need to do this because Replicate returns a list of strings (for streaming functionality which is not exposed by the run function)
|
||||
role: "assistant",
|
||||
},
|
||||
};
|
||||
} as R;
|
||||
}
|
||||
|
||||
async complete(
|
||||
prompt: string,
|
||||
parentEvent?: Event,
|
||||
): Promise<CompletionResponse> {
|
||||
async complete<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(prompt: string, parentEvent?: Event, streaming?: T): Promise<R> {
|
||||
return this.chat([{ content: prompt, role: "user" }], parentEvent);
|
||||
}
|
||||
}
|
||||
|
||||
export const ALL_AVAILABLE_ANTHROPIC_MODELS = {
|
||||
// both models have 100k context window, see https://docs.anthropic.com/claude/reference/selecting-a-model
|
||||
"claude-2": { contextWindow: 200000 },
|
||||
"claude-instant-1": { contextWindow: 100000 },
|
||||
};
|
||||
|
||||
/**
|
||||
* Anthropic LLM implementation
|
||||
*/
|
||||
|
||||
export class Anthropic implements LLM {
|
||||
hasStreaming: boolean = true;
|
||||
|
||||
// Per completion Anthropic params
|
||||
model: string;
|
||||
model: keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS;
|
||||
temperature: number;
|
||||
topP: number;
|
||||
maxTokens?: number;
|
||||
@@ -483,25 +684,55 @@ export class Anthropic implements LLM {
|
||||
this.callbackManager = init?.callbackManager;
|
||||
}
|
||||
|
||||
tokens(messages: ChatMessage[]): number {
|
||||
throw new Error("Method not implemented.");
|
||||
}
|
||||
|
||||
get metadata() {
|
||||
return {
|
||||
model: this.model,
|
||||
temperature: this.temperature,
|
||||
topP: this.topP,
|
||||
maxTokens: this.maxTokens,
|
||||
contextWindow: ALL_AVAILABLE_ANTHROPIC_MODELS[this.model].contextWindow,
|
||||
tokenizer: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
mapMessagesToPrompt(messages: ChatMessage[]) {
|
||||
return (
|
||||
messages.reduce((acc, message) => {
|
||||
return (
|
||||
acc +
|
||||
`${
|
||||
message.role === "assistant"
|
||||
? ANTHROPIC_AI_PROMPT
|
||||
: ANTHROPIC_HUMAN_PROMPT
|
||||
} ${message.content} `
|
||||
message.role === "system"
|
||||
? ""
|
||||
: message.role === "assistant"
|
||||
? ANTHROPIC_AI_PROMPT + " "
|
||||
: ANTHROPIC_HUMAN_PROMPT + " "
|
||||
}${message.content.trim()}`
|
||||
);
|
||||
}, "") + ANTHROPIC_AI_PROMPT
|
||||
);
|
||||
}
|
||||
|
||||
async chat(
|
||||
async chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(
|
||||
messages: ChatMessage[],
|
||||
parentEvent?: Event | undefined,
|
||||
): Promise<ChatResponse> {
|
||||
streaming?: T,
|
||||
): Promise<R> {
|
||||
//Streaming
|
||||
if (streaming) {
|
||||
if (!this.hasStreaming) {
|
||||
throw Error("No streaming support for this LLM.");
|
||||
}
|
||||
return this.streamChat(messages, parentEvent) as R;
|
||||
}
|
||||
|
||||
//Non-streaming
|
||||
const response = await this.session.anthropic.completions.create({
|
||||
model: this.model,
|
||||
prompt: this.mapMessagesToPrompt(messages),
|
||||
@@ -514,12 +745,182 @@ export class Anthropic implements LLM {
|
||||
message: { content: response.completion.trimStart(), role: "assistant" },
|
||||
//^ We're trimming the start because Anthropic often starts with a space in the response
|
||||
// That space will be re-added when we generate the next prompt.
|
||||
};
|
||||
} as R;
|
||||
}
|
||||
async complete(
|
||||
|
||||
protected async *streamChat(
|
||||
messages: ChatMessage[],
|
||||
parentEvent?: Event | undefined,
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
// AsyncIterable<AnthropicStreamToken>
|
||||
const stream: AsyncIterable<AnthropicStreamToken> =
|
||||
await this.session.anthropic.completions.create({
|
||||
model: this.model,
|
||||
prompt: this.mapMessagesToPrompt(messages),
|
||||
max_tokens_to_sample: this.maxTokens ?? 100000,
|
||||
temperature: this.temperature,
|
||||
top_p: this.topP,
|
||||
stream: true,
|
||||
});
|
||||
|
||||
var idx_counter: number = 0;
|
||||
for await (const part of stream) {
|
||||
//TODO: LLM Stream Callback, pending re-work.
|
||||
|
||||
idx_counter++;
|
||||
yield part.completion;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
async complete<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(
|
||||
prompt: string,
|
||||
parentEvent?: Event | undefined,
|
||||
): Promise<CompletionResponse> {
|
||||
return this.chat([{ content: prompt, role: "user" }], parentEvent);
|
||||
streaming?: T,
|
||||
): Promise<R> {
|
||||
if (streaming) {
|
||||
return this.streamComplete(prompt, parentEvent) as R;
|
||||
}
|
||||
return this.chat(
|
||||
[{ content: prompt, role: "user" }],
|
||||
parentEvent,
|
||||
streaming,
|
||||
) as R;
|
||||
}
|
||||
|
||||
protected streamComplete(
|
||||
prompt: string,
|
||||
parentEvent?: Event | undefined,
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
return this.streamChat([{ content: prompt, role: "user" }], parentEvent);
|
||||
}
|
||||
}
|
||||
|
||||
export class Portkey implements LLM {
|
||||
hasStreaming: boolean = true;
|
||||
|
||||
apiKey?: string = undefined;
|
||||
baseURL?: string = undefined;
|
||||
mode?: string = undefined;
|
||||
llms?: [LLMOptions] | null = undefined;
|
||||
session: PortkeySession;
|
||||
callbackManager?: CallbackManager;
|
||||
|
||||
constructor(init?: Partial<Portkey>) {
|
||||
this.apiKey = init?.apiKey;
|
||||
this.baseURL = init?.baseURL;
|
||||
this.mode = init?.mode;
|
||||
this.llms = init?.llms;
|
||||
this.session = getPortkeySession({
|
||||
apiKey: this.apiKey,
|
||||
baseURL: this.baseURL,
|
||||
llms: this.llms,
|
||||
mode: this.mode,
|
||||
});
|
||||
this.callbackManager = init?.callbackManager;
|
||||
}
|
||||
|
||||
tokens(messages: ChatMessage[]): number {
|
||||
throw new Error("Method not implemented.");
|
||||
}
|
||||
|
||||
get metadata(): LLMMetadata {
|
||||
throw new Error("metadata not implemented for Portkey");
|
||||
}
|
||||
|
||||
async chat<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(
|
||||
messages: ChatMessage[],
|
||||
parentEvent?: Event | undefined,
|
||||
streaming?: T,
|
||||
params?: Record<string, any>,
|
||||
): Promise<R> {
|
||||
if (streaming) {
|
||||
return this.streamChat(messages, parentEvent, params) as R;
|
||||
} else {
|
||||
const resolvedParams = params || {};
|
||||
const response = await this.session.portkey.chatCompletions.create({
|
||||
messages,
|
||||
...resolvedParams,
|
||||
});
|
||||
|
||||
const content = response.choices[0].message?.content ?? "";
|
||||
const role = response.choices[0].message?.role || "assistant";
|
||||
return { message: { content, role: role as MessageType } } as R;
|
||||
}
|
||||
}
|
||||
|
||||
async complete<
|
||||
T extends boolean | undefined = undefined,
|
||||
R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
|
||||
>(
|
||||
prompt: string,
|
||||
parentEvent?: Event | undefined,
|
||||
streaming?: T,
|
||||
): Promise<R> {
|
||||
return this.chat(
|
||||
[{ content: prompt, role: "user" }],
|
||||
parentEvent,
|
||||
streaming,
|
||||
);
|
||||
}
|
||||
|
||||
async *streamChat(
|
||||
messages: ChatMessage[],
|
||||
parentEvent?: Event,
|
||||
params?: Record<string, any>,
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
// Wrapping the stream in a callback.
|
||||
const onLLMStream = this.callbackManager?.onLLMStream
|
||||
? this.callbackManager.onLLMStream
|
||||
: () => {};
|
||||
|
||||
const chunkStream = await this.session.portkey.chatCompletions.create({
|
||||
messages,
|
||||
...params,
|
||||
stream: true,
|
||||
});
|
||||
|
||||
const event: Event = parentEvent
|
||||
? parentEvent
|
||||
: {
|
||||
id: "unspecified",
|
||||
type: "llmPredict" as EventType,
|
||||
};
|
||||
|
||||
//Indices
|
||||
var idx_counter: number = 0;
|
||||
for await (const part of chunkStream) {
|
||||
//Increment
|
||||
part.choices[0].index = idx_counter;
|
||||
const is_done: boolean =
|
||||
part.choices[0].finish_reason === "stop" ? true : false;
|
||||
//onLLMStream Callback
|
||||
|
||||
const stream_callback: StreamCallbackResponse = {
|
||||
event: event,
|
||||
index: idx_counter,
|
||||
isDone: is_done,
|
||||
// token: part,
|
||||
};
|
||||
onLLMStream(stream_callback);
|
||||
|
||||
idx_counter++;
|
||||
|
||||
yield part.choices[0].delta?.content ?? "";
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
streamComplete(
|
||||
query: string,
|
||||
parentEvent?: Event,
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
return this.streamChat([{ content: query, role: "user" }], parentEvent);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -24,7 +24,10 @@ export class OpenAISession {
|
||||
if (options.azure) {
|
||||
this.openai = new AzureOpenAI(options);
|
||||
} else {
|
||||
this.openai = new OpenAI(options);
|
||||
this.openai = new OpenAI({
|
||||
...options,
|
||||
// defaultHeaders: { "OpenAI-Beta": "assistants=v1" },
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
import _ from "lodash";
|
||||
import { LLMOptions, Portkey } from "portkey-ai";
|
||||
|
||||
export const readEnv = (
|
||||
env: string,
|
||||
default_val?: string,
|
||||
): string | undefined => {
|
||||
if (typeof process !== "undefined") {
|
||||
return process.env?.[env] ?? default_val;
|
||||
}
|
||||
return default_val;
|
||||
};
|
||||
|
||||
interface PortkeyOptions {
|
||||
apiKey?: string;
|
||||
baseURL?: string;
|
||||
mode?: string;
|
||||
llms?: [LLMOptions] | null;
|
||||
}
|
||||
|
||||
export class PortkeySession {
|
||||
portkey: Portkey;
|
||||
|
||||
constructor(options: PortkeyOptions = {}) {
|
||||
if (!options.apiKey) {
|
||||
options.apiKey = readEnv("PORTKEY_API_KEY");
|
||||
}
|
||||
|
||||
if (!options.baseURL) {
|
||||
options.baseURL = readEnv("PORTKEY_BASE_URL", "https://api.portkey.ai");
|
||||
}
|
||||
|
||||
this.portkey = new Portkey({});
|
||||
this.portkey.llms = [{}];
|
||||
if (!options.apiKey) {
|
||||
throw new Error("Set Portkey ApiKey in PORTKEY_API_KEY env variable");
|
||||
}
|
||||
|
||||
this.portkey = new Portkey(options);
|
||||
}
|
||||
}
|
||||
|
||||
let defaultPortkeySession: {
|
||||
session: PortkeySession;
|
||||
options: PortkeyOptions;
|
||||
}[] = [];
|
||||
|
||||
/**
|
||||
* Get a session for the Portkey API. If one already exists with the same options,
|
||||
* it will be returned. Otherwise, a new session will be created.
|
||||
* @param options
|
||||
* @returns
|
||||
*/
|
||||
export function getPortkeySession(options: PortkeyOptions = {}) {
|
||||
let session = defaultPortkeySession.find((session) => {
|
||||
return _.isEqual(session.options, options);
|
||||
})?.session;
|
||||
|
||||
if (!session) {
|
||||
session = new PortkeySession(options);
|
||||
defaultPortkeySession.push({ session, options });
|
||||
}
|
||||
return session;
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
import mammoth from "mammoth";
|
||||
import { Document } from "../Node";
|
||||
import { DEFAULT_FS } from "../storage/constants";
|
||||
import { GenericFileSystem } from "../storage/FileSystem";
|
||||
import { BaseReader } from "./base";
|
||||
|
||||
export class DocxReader implements BaseReader {
|
||||
/** DocxParser */
|
||||
async loadData(
|
||||
file: string,
|
||||
fs: GenericFileSystem = DEFAULT_FS,
|
||||
): Promise<Document[]> {
|
||||
const dataBuffer = (await fs.readFile(file)) as any;
|
||||
const { value } = await mammoth.extractRawText({ buffer: dataBuffer });
|
||||
return [new Document({ text: value, id_: file })];
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
import { Document } from "../Node";
|
||||
import { DEFAULT_FS } from "../storage/constants";
|
||||
import { GenericFileSystem } from "../storage/FileSystem";
|
||||
import { BaseReader } from "./base";
|
||||
|
||||
/**
|
||||
* Extract the significant text from an arbitrary HTML document.
|
||||
* The contents of any head, script, style, and xml tags are removed completely.
|
||||
* The URLs for a[href] tags are extracted, along with the inner text of the tag.
|
||||
* All other tags are removed, and the inner text is kept intact.
|
||||
* Html entities (e.g., &) are not decoded.
|
||||
*/
|
||||
export class HTMLReader implements BaseReader {
|
||||
/**
|
||||
* Public method for this reader.
|
||||
* Required by BaseReader interface.
|
||||
* @param file Path/name of the file to be loaded.
|
||||
* @param fs fs wrapper interface for getting the file content.
|
||||
* @returns Promise<Document[]> A Promise object, eventually yielding zero or one Document parsed from the HTML content of the specified file.
|
||||
*/
|
||||
async loadData(
|
||||
file: string,
|
||||
fs: GenericFileSystem = DEFAULT_FS,
|
||||
): Promise<Document[]> {
|
||||
const dataBuffer = await fs.readFile(file, "utf-8");
|
||||
const htmlOptions = this.getOptions();
|
||||
const content = await this.parseContent(dataBuffer, htmlOptions);
|
||||
return [new Document({ text: content, id_: file })];
|
||||
}
|
||||
|
||||
/**
|
||||
* Wrapper for string-strip-html usage.
|
||||
* @param html Raw HTML content to be parsed.
|
||||
* @param options An object of options for the underlying library
|
||||
* @see getOptions
|
||||
* @returns The HTML content, stripped of unwanted tags and attributes
|
||||
*/
|
||||
async parseContent(html: string, options: any = {}): Promise<string> {
|
||||
const { stripHtml } = await import("string-strip-html"); // ESM only
|
||||
return stripHtml(html).result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Wrapper for our configuration options passed to string-strip-html library
|
||||
* @see https://codsen.com/os/string-strip-html/examples
|
||||
* @returns An object of options for the underlying library
|
||||
*/
|
||||
getOptions() {
|
||||
return {
|
||||
skipHtmlDecoding: true,
|
||||
stripTogetherWithTheirContents: [
|
||||
"script", // default
|
||||
"style", // default
|
||||
"xml", // default
|
||||
"head", // <-- custom-added
|
||||
],
|
||||
// Keep the URLs for embedded links
|
||||
// cb: (tag: any, deleteFrom: number, deleteTo: number, insert: string, rangesArr: any, proposedReturn: string) => {
|
||||
// let temp;
|
||||
// if (
|
||||
// tag.name === "a" &&
|
||||
// tag.attributes &&
|
||||
// tag.attributes.some((attr: any) => {
|
||||
// if (attr.name === "href") {
|
||||
// temp = attr.value;
|
||||
// return true;
|
||||
// }
|
||||
// })
|
||||
// ) {
|
||||
// rangesArr.push([deleteFrom, deleteTo, `${temp} ${insert || ""}`]);
|
||||
// } else {
|
||||
// rangesArr.push(proposedReturn);
|
||||
// }
|
||||
// },
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -3,10 +3,24 @@ import { Document } from "../Node";
|
||||
import { CompleteFileSystem, walk } from "../storage/FileSystem";
|
||||
import { DEFAULT_FS } from "../storage/constants";
|
||||
import { PapaCSVReader } from "./CSVReader";
|
||||
import { DocxReader } from "./DocxReader";
|
||||
import { HTMLReader } from "./HTMLReader";
|
||||
import { MarkdownReader } from "./MarkdownReader";
|
||||
import { PDFReader } from "./PDFReader";
|
||||
import { BaseReader } from "./base";
|
||||
|
||||
type ReaderCallback = (
|
||||
category: "file" | "directory",
|
||||
name: string,
|
||||
status: ReaderStatus,
|
||||
message?: string,
|
||||
) => boolean;
|
||||
enum ReaderStatus {
|
||||
STARTED = 0,
|
||||
COMPLETE,
|
||||
ERROR,
|
||||
}
|
||||
|
||||
/**
|
||||
* Read a .txt file
|
||||
*/
|
||||
@@ -20,11 +34,14 @@ export class TextFileReader implements BaseReader {
|
||||
}
|
||||
}
|
||||
|
||||
const FILE_EXT_TO_READER: Record<string, BaseReader> = {
|
||||
export const FILE_EXT_TO_READER: Record<string, BaseReader> = {
|
||||
txt: new TextFileReader(),
|
||||
pdf: new PDFReader(),
|
||||
csv: new PapaCSVReader(),
|
||||
md: new MarkdownReader(),
|
||||
docx: new DocxReader(),
|
||||
htm: new HTMLReader(),
|
||||
html: new HTMLReader(),
|
||||
};
|
||||
|
||||
export type SimpleDirectoryReaderLoadDataProps = {
|
||||
@@ -35,20 +52,37 @@ export type SimpleDirectoryReaderLoadDataProps = {
|
||||
};
|
||||
|
||||
/**
|
||||
* Read all of the documents in a directory. Currently supports PDF and TXT files.
|
||||
* Read all of the documents in a directory.
|
||||
* By default, supports the list of file types
|
||||
* in the FILE_EXIT_TO_READER map.
|
||||
*/
|
||||
export class SimpleDirectoryReader implements BaseReader {
|
||||
constructor(private observer?: ReaderCallback) {}
|
||||
|
||||
async loadData({
|
||||
directoryPath,
|
||||
fs = DEFAULT_FS as CompleteFileSystem,
|
||||
defaultReader = new TextFileReader(),
|
||||
fileExtToReader = FILE_EXT_TO_READER,
|
||||
}: SimpleDirectoryReaderLoadDataProps): Promise<Document[]> {
|
||||
// Observer can decide to skip the directory
|
||||
if (
|
||||
!this.doObserverCheck("directory", directoryPath, ReaderStatus.STARTED)
|
||||
) {
|
||||
return [];
|
||||
}
|
||||
|
||||
let docs: Document[] = [];
|
||||
for await (const filePath of walk(fs, directoryPath)) {
|
||||
try {
|
||||
const fileExt = _.last(filePath.split(".")) || "";
|
||||
|
||||
// Observer can decide to skip each file
|
||||
if (!this.doObserverCheck("file", filePath, ReaderStatus.STARTED)) {
|
||||
// Skip this file
|
||||
continue;
|
||||
}
|
||||
|
||||
let reader = null;
|
||||
|
||||
if (fileExt in fileExtToReader) {
|
||||
@@ -56,16 +90,52 @@ export class SimpleDirectoryReader implements BaseReader {
|
||||
} else if (!_.isNil(defaultReader)) {
|
||||
reader = defaultReader;
|
||||
} else {
|
||||
console.warn(`No reader for file extension of ${filePath}`);
|
||||
const msg = `No reader for file extension of ${filePath}`;
|
||||
console.warn(msg);
|
||||
|
||||
// In an error condition, observer's false cancels the whole process.
|
||||
if (
|
||||
!this.doObserverCheck("file", filePath, ReaderStatus.ERROR, msg)
|
||||
) {
|
||||
return [];
|
||||
}
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
const fileDocs = await reader.loadData(filePath, fs);
|
||||
docs.push(...fileDocs);
|
||||
|
||||
// Observer can still cancel addition of the resulting docs from this file
|
||||
if (this.doObserverCheck("file", filePath, ReaderStatus.COMPLETE)) {
|
||||
docs.push(...fileDocs);
|
||||
}
|
||||
} catch (e) {
|
||||
console.error(`Error reading file ${filePath}: ${e}`);
|
||||
const msg = `Error reading file ${filePath}: ${e}`;
|
||||
console.error(msg);
|
||||
|
||||
// In an error condition, observer's false cancels the whole process.
|
||||
if (!this.doObserverCheck("file", filePath, ReaderStatus.ERROR, msg)) {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// After successful import of all files, directory completion
|
||||
// is only a notification for observer, cannot be cancelled.
|
||||
this.doObserverCheck("directory", directoryPath, ReaderStatus.COMPLETE);
|
||||
|
||||
return docs;
|
||||
}
|
||||
|
||||
private doObserverCheck(
|
||||
category: "file" | "directory",
|
||||
name: string,
|
||||
status: ReaderStatus,
|
||||
message?: string,
|
||||
): boolean {
|
||||
if (this.observer) {
|
||||
return this.observer(category, name, status, message);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
import { MongoClient } from "mongodb";
|
||||
import { Document, Metadata } from "../Node";
|
||||
import { BaseReader } from "./base";
|
||||
|
||||
/**
|
||||
* Read in from MongoDB
|
||||
*/
|
||||
export class SimpleMongoReader implements BaseReader {
|
||||
private client: MongoClient;
|
||||
|
||||
constructor(client: MongoClient) {
|
||||
this.client = client;
|
||||
}
|
||||
|
||||
/**
|
||||
* Flattens an array of strings or string arrays into a single-dimensional array of strings.
|
||||
* @param texts - The array of strings or string arrays to flatten.
|
||||
* @returns The flattened array of strings.
|
||||
*/
|
||||
private flatten(texts: Array<string | string[]>): string[] {
|
||||
return texts.reduce<string[]>(
|
||||
(result, text) => result.concat(text instanceof Array ? text : [text]),
|
||||
[],
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads data from MongoDB collection
|
||||
* @param {string} dbName - The name of the database to load.
|
||||
* @param {string} collectionName - The name of the collection to load.
|
||||
* @param {string[]} fieldNames - An array of field names to retrieve from each document. Defaults to ["text"].
|
||||
* @param {string} separator - The separator to join multiple field values. Defaults to an empty string.
|
||||
* @param {Record<string, any>} filterQuery - Specific query, as specified by MongoDB NodeJS documentation.
|
||||
* @param {Number} maxDocs - The maximum number of documents to retrieve. Defaults to 0 (retrieve all documents).
|
||||
* @param {string[]} metadataNames - An optional array of metadata field names. If specified extracts this information as metadata.
|
||||
* @returns {Promise<Document[]>}
|
||||
* @throws If a field specified in fieldNames or metadataNames is not found in a MongoDB document.
|
||||
*/
|
||||
public async loadData(
|
||||
dbName: string,
|
||||
collectionName: string,
|
||||
fieldNames: string[] = ["text"],
|
||||
separator: string = "",
|
||||
filterQuery: Record<string, any> = {},
|
||||
maxDocs: number = 0,
|
||||
metadataNames?: string[],
|
||||
): Promise<Document[]> {
|
||||
const db = this.client.db(dbName);
|
||||
// Get items from collection
|
||||
const cursor = db
|
||||
.collection(collectionName)
|
||||
.find(filterQuery)
|
||||
.limit(maxDocs);
|
||||
|
||||
const documents: Document[] = [];
|
||||
|
||||
for await (const item of cursor) {
|
||||
try {
|
||||
const texts: Array<string | string[]> = fieldNames.map(
|
||||
(name) => item[name],
|
||||
);
|
||||
const flattenedTexts = this.flatten(texts);
|
||||
const text = flattenedTexts.join(separator);
|
||||
|
||||
let metadata: Metadata = {};
|
||||
if (metadataNames) {
|
||||
// extract metadata if fields are specified
|
||||
metadata = Object.fromEntries(
|
||||
metadataNames.map((name) => [name, item[name]]),
|
||||
);
|
||||
}
|
||||
|
||||
documents.push(new Document({ text, metadata }));
|
||||
} catch (err) {
|
||||
throw new Error(
|
||||
`Field not found in Mongo document: ${(err as Error).message}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
return documents;
|
||||
}
|
||||
}
|
||||
@@ -7,5 +7,6 @@ export { SimpleIndexStore } from "./indexStore/SimpleIndexStore";
|
||||
export * from "./indexStore/types";
|
||||
export { SimpleKVStore } from "./kvStore/SimpleKVStore";
|
||||
export * from "./kvStore/types";
|
||||
export { MongoDBAtlasVectorSearch } from "./vectorStore/MongoDBAtlasVectorStore";
|
||||
export { SimpleVectorStore } from "./vectorStore/SimpleVectorStore";
|
||||
export * from "./vectorStore/types";
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
import { BulkWriteOptions, Collection, MongoClient } from "mongodb";
|
||||
import { BaseNode, MetadataMode } from "../../Node";
|
||||
import {
|
||||
MetadataFilters,
|
||||
VectorStore,
|
||||
VectorStoreQuery,
|
||||
VectorStoreQueryResult,
|
||||
} from "./types";
|
||||
import { metadataDictToNode, nodeToMetadata } from "./utils";
|
||||
|
||||
// Utility function to convert metadata filters to MongoDB filter
|
||||
function toMongoDBFilter(
|
||||
standardFilters: MetadataFilters,
|
||||
): Record<string, any> {
|
||||
const filters: Record<string, any> = {};
|
||||
for (const filter of standardFilters.filters) {
|
||||
filters[filter.key] = filter.value;
|
||||
}
|
||||
return filters;
|
||||
}
|
||||
|
||||
// MongoDB Atlas Vector Store class implementing VectorStore
|
||||
export class MongoDBAtlasVectorSearch implements VectorStore {
|
||||
storesText: boolean = true;
|
||||
flatMetadata: boolean = true;
|
||||
|
||||
mongodbClient: MongoClient;
|
||||
indexName: string;
|
||||
embeddingKey: string;
|
||||
idKey: string;
|
||||
textKey: string;
|
||||
metadataKey: string;
|
||||
insertOptions?: BulkWriteOptions;
|
||||
private collection: Collection;
|
||||
|
||||
constructor(
|
||||
init: Partial<MongoDBAtlasVectorSearch> & {
|
||||
dbName: string;
|
||||
collectionName: string;
|
||||
},
|
||||
) {
|
||||
if (init.mongodbClient) {
|
||||
this.mongodbClient = init.mongodbClient;
|
||||
} else {
|
||||
const mongoUri = process.env.MONGODB_URI;
|
||||
if (!mongoUri) {
|
||||
throw new Error(
|
||||
"Must specify MONGODB_URI via env variable if not directly passing in client.",
|
||||
);
|
||||
}
|
||||
this.mongodbClient = new MongoClient(mongoUri);
|
||||
}
|
||||
|
||||
this.collection = this.mongodbClient
|
||||
.db(init.dbName ?? "default_db")
|
||||
.collection(init.collectionName ?? "default_collection");
|
||||
this.indexName = init.indexName ?? "default";
|
||||
this.embeddingKey = init.embeddingKey ?? "embedding";
|
||||
this.idKey = init.idKey ?? "id";
|
||||
this.textKey = init.textKey ?? "text";
|
||||
this.metadataKey = init.metadataKey ?? "metadata";
|
||||
this.insertOptions = init.insertOptions;
|
||||
}
|
||||
|
||||
async add(nodes: BaseNode[]): Promise<string[]> {
|
||||
if (!nodes || nodes.length === 0) {
|
||||
return [];
|
||||
}
|
||||
const dataToInsert = nodes.map((node) => {
|
||||
const metadata = nodeToMetadata(
|
||||
node,
|
||||
true,
|
||||
this.textKey,
|
||||
this.flatMetadata,
|
||||
);
|
||||
|
||||
return {
|
||||
[this.idKey]: node.id_,
|
||||
[this.embeddingKey]: node.getEmbedding(),
|
||||
[this.textKey]: node.getContent(MetadataMode.NONE) || "",
|
||||
[this.metadataKey]: metadata,
|
||||
};
|
||||
});
|
||||
|
||||
console.debug("Inserting data into MongoDB: ", dataToInsert);
|
||||
const insertResult = await this.collection.insertMany(
|
||||
dataToInsert,
|
||||
this.insertOptions,
|
||||
);
|
||||
console.debug("Result of insert: ", insertResult);
|
||||
return nodes.map((node) => node.id_);
|
||||
}
|
||||
|
||||
async delete(refDocId: string, deleteOptions?: any): Promise<void> {
|
||||
await this.collection.deleteOne(
|
||||
{
|
||||
[`${this.metadataKey}.ref_doc_id`]: refDocId,
|
||||
},
|
||||
deleteOptions,
|
||||
);
|
||||
}
|
||||
|
||||
get client(): any {
|
||||
return this.mongodbClient;
|
||||
}
|
||||
|
||||
async query(
|
||||
query: VectorStoreQuery,
|
||||
options?: any,
|
||||
): Promise<VectorStoreQueryResult> {
|
||||
const params: any = {
|
||||
queryVector: query.queryEmbedding,
|
||||
path: this.embeddingKey,
|
||||
numCandidates: query.similarityTopK * 10,
|
||||
limit: query.similarityTopK,
|
||||
index: this.indexName,
|
||||
};
|
||||
|
||||
if (query.filters) {
|
||||
params.filter = toMongoDBFilter(query.filters);
|
||||
}
|
||||
|
||||
const queryField = { $vectorSearch: params };
|
||||
const pipeline = [
|
||||
queryField,
|
||||
{
|
||||
$project: {
|
||||
score: { $meta: "vectorSearchScore" },
|
||||
[this.embeddingKey]: 0,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
console.debug("Running query pipeline: ", pipeline);
|
||||
const cursor = await this.collection.aggregate(pipeline);
|
||||
|
||||
const nodes: BaseNode[] = [];
|
||||
const ids: string[] = [];
|
||||
const similarities: number[] = [];
|
||||
|
||||
for await (const res of await cursor) {
|
||||
const text = res[this.textKey];
|
||||
const score = res.score;
|
||||
const id = res[this.idKey];
|
||||
const metadata = res[this.metadataKey];
|
||||
|
||||
const node = metadataDictToNode(metadata);
|
||||
node.setContent(text);
|
||||
|
||||
ids.push(id);
|
||||
nodes.push(node);
|
||||
similarities.push(score);
|
||||
}
|
||||
|
||||
const result = {
|
||||
nodes,
|
||||
similarities,
|
||||
ids,
|
||||
};
|
||||
|
||||
console.debug("Result of query (ids):", ids);
|
||||
return result;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,266 @@
|
||||
import pg from "pg";
|
||||
import pgvector from "pgvector/pg";
|
||||
|
||||
import { VectorStore, VectorStoreQuery, VectorStoreQueryResult } from "./types";
|
||||
|
||||
import { BaseNode, Document, Metadata, MetadataMode } from "../../Node";
|
||||
import { GenericFileSystem } from "../FileSystem";
|
||||
|
||||
export const PGVECTOR_SCHEMA = "public";
|
||||
export const PGVECTOR_TABLE = "llamaindex_embedding";
|
||||
|
||||
/**
|
||||
* Provides support for writing and querying vector data in Postgres.
|
||||
*/
|
||||
export class PGVectorStore implements VectorStore {
|
||||
storesText: boolean = true;
|
||||
|
||||
private collection: string = "";
|
||||
|
||||
/*
|
||||
FROM pg LIBRARY:
|
||||
type Config = {
|
||||
user?: string, // default process.env.PGUSER || process.env.USER
|
||||
password?: string or function, //default process.env.PGPASSWORD
|
||||
host?: string, // default process.env.PGHOST
|
||||
database?: string, // default process.env.PGDATABASE || user
|
||||
port?: number, // default process.env.PGPORT
|
||||
connectionString?: string, // e.g. postgres://user:password@host:5432/database
|
||||
ssl?: any, // passed directly to node.TLSSocket, supports all tls.connect options
|
||||
types?: any, // custom type parsers
|
||||
statement_timeout?: number, // number of milliseconds before a statement in query will time out, default is no timeout
|
||||
query_timeout?: number, // number of milliseconds before a query call will timeout, default is no timeout
|
||||
application_name?: string, // The name of the application that created this Client instance
|
||||
connectionTimeoutMillis?: number, // number of milliseconds to wait for connection, default is no timeout
|
||||
idle_in_transaction_session_timeout?: number // number of milliseconds before terminating any session with an open idle transaction, default is no timeout
|
||||
}
|
||||
*/
|
||||
db?: pg.Client;
|
||||
|
||||
constructor() {}
|
||||
|
||||
/**
|
||||
* Setter for the collection property.
|
||||
* Using a collection allows for simple segregation of vector data,
|
||||
* e.g. by user, source, or access-level.
|
||||
* Leave/set blank to ignore the collection value when querying.
|
||||
* @param coll Name for the collection.
|
||||
*/
|
||||
setCollection(coll: string) {
|
||||
this.collection = coll;
|
||||
}
|
||||
|
||||
/**
|
||||
* Getter for the collection property.
|
||||
* Using a collection allows for simple segregation of vector data,
|
||||
* e.g. by user, source, or access-level.
|
||||
* Leave/set blank to ignore the collection value when querying.
|
||||
* @returns The currently-set collection value. Default is empty string.
|
||||
*/
|
||||
getCollection(): string {
|
||||
return this.collection;
|
||||
}
|
||||
|
||||
private async getDb(): Promise<pg.Client> {
|
||||
if (!this.db) {
|
||||
try {
|
||||
// Create DB connection
|
||||
// Read connection params from env - see comment block above
|
||||
const db = new pg.Client();
|
||||
await db.connect();
|
||||
|
||||
// Check vector extension
|
||||
db.query("CREATE EXTENSION IF NOT EXISTS vector");
|
||||
await pgvector.registerType(db);
|
||||
|
||||
// Check schema, table(s), index(es)
|
||||
await this.checkSchema(db);
|
||||
|
||||
// All good? Keep the connection reference
|
||||
this.db = db;
|
||||
} catch (err: any) {
|
||||
console.error(err);
|
||||
return Promise.reject(err);
|
||||
}
|
||||
}
|
||||
|
||||
return Promise.resolve(this.db);
|
||||
}
|
||||
|
||||
private async checkSchema(db: pg.Client) {
|
||||
await db.query(`CREATE SCHEMA IF NOT EXISTS ${PGVECTOR_SCHEMA}`);
|
||||
|
||||
const tbl = `CREATE TABLE IF NOT EXISTS ${PGVECTOR_SCHEMA}.${PGVECTOR_TABLE}(
|
||||
id uuid DEFAULT gen_random_uuid() PRIMARY KEY,
|
||||
external_id VARCHAR,
|
||||
collection VARCHAR,
|
||||
document TEXT,
|
||||
metadata JSONB DEFAULT '{}',
|
||||
embeddings VECTOR(1536)
|
||||
)`;
|
||||
await db.query(tbl);
|
||||
|
||||
const idxs = `CREATE INDEX IF NOT EXISTS idx_${PGVECTOR_TABLE}_external_id ON ${PGVECTOR_SCHEMA}.${PGVECTOR_TABLE} (external_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_${PGVECTOR_TABLE}_collection ON ${PGVECTOR_SCHEMA}.${PGVECTOR_TABLE} (collection);`;
|
||||
await db.query(idxs);
|
||||
|
||||
// TODO add IVFFlat or HNSW indexing?
|
||||
return db;
|
||||
}
|
||||
|
||||
// isEmbeddingQuery?: boolean | undefined;
|
||||
|
||||
/**
|
||||
* Connects to the database specified in environment vars.
|
||||
* This method also checks and creates the vector extension,
|
||||
* the destination table and indexes if not found.
|
||||
* @returns A connection to the database, or the error encountered while connecting/setting up.
|
||||
*/
|
||||
client() {
|
||||
return this.getDb();
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete all vector records for the specified collection.
|
||||
* NOTE: Uses the collection property controlled by setCollection/getCollection.
|
||||
* @returns The result of the delete query.
|
||||
*/
|
||||
async clearCollection() {
|
||||
const sql: string = `DELETE FROM ${PGVECTOR_SCHEMA}.${PGVECTOR_TABLE}
|
||||
WHERE collection = $1`;
|
||||
|
||||
const db = (await this.getDb()) as pg.Client;
|
||||
const ret = await db.query(sql, [this.collection]);
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds vector record(s) to the table.
|
||||
* NOTE: Uses the collection property controlled by setCollection/getCollection.
|
||||
* @param embeddingResults The Nodes to be inserted, optionally including metadata tuples.
|
||||
* @returns A list of zero or more id values for the created records.
|
||||
*/
|
||||
async add(embeddingResults: BaseNode<Metadata>[]): Promise<string[]> {
|
||||
const sql: string = `INSERT INTO ${PGVECTOR_SCHEMA}.${PGVECTOR_TABLE}
|
||||
(id, external_id, collection, document, metadata, embeddings)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)`;
|
||||
|
||||
const db = (await this.getDb()) as pg.Client;
|
||||
|
||||
let ret: string[] = [];
|
||||
for (let index = 0; index < embeddingResults.length; index++) {
|
||||
const row = embeddingResults[index];
|
||||
|
||||
let id: any = row.id_.length ? row.id_ : null;
|
||||
let meta = row.metadata || {};
|
||||
meta.create_date = new Date();
|
||||
|
||||
const params = [
|
||||
id,
|
||||
"",
|
||||
this.collection,
|
||||
row.getContent(MetadataMode.EMBED),
|
||||
meta,
|
||||
"[" + row.getEmbedding().join(",") + "]",
|
||||
];
|
||||
|
||||
try {
|
||||
const result = await db.query(sql, params);
|
||||
|
||||
if (result.rows.length) {
|
||||
id = result.rows[0].id as string;
|
||||
ret.push(id);
|
||||
}
|
||||
} catch (err) {
|
||||
const msg = `${err}`;
|
||||
console.log(msg, err);
|
||||
}
|
||||
}
|
||||
|
||||
return Promise.resolve(ret);
|
||||
}
|
||||
|
||||
/**
|
||||
* 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 delete(refDocId: string, deleteKwargs?: any): Promise<void> {
|
||||
const collectionCriteria = this.collection.length
|
||||
? "AND collection = $2"
|
||||
: "";
|
||||
const sql: string = `DELETE FROM ${PGVECTOR_SCHEMA}.${PGVECTOR_TABLE}
|
||||
WHERE id = $1 ${collectionCriteria}`;
|
||||
|
||||
const db = (await this.getDb()) as pg.Client;
|
||||
const params = this.collection.length
|
||||
? [refDocId, this.collection]
|
||||
: [refDocId];
|
||||
await db.query(sql, params);
|
||||
return Promise.resolve();
|
||||
}
|
||||
|
||||
/**
|
||||
* 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> {
|
||||
// TODO QUERY TYPES:
|
||||
// Distance: SELECT embedding <-> $1 AS distance FROM items;
|
||||
// Inner Product: SELECT (embedding <#> $1) * -1 AS inner_product FROM items;
|
||||
// Cosine Sim: SELECT 1 - (embedding <=> $1) AS cosine_similarity FROM items;
|
||||
|
||||
const embedding = "[" + query.queryEmbedding?.join(",") + "]";
|
||||
const max = query.similarityTopK ?? 2;
|
||||
const where = this.collection.length ? "WHERE collection = $2" : "";
|
||||
// TODO Add collection filter if set
|
||||
const sql = `SELECT * FROM ${PGVECTOR_SCHEMA}.${PGVECTOR_TABLE}
|
||||
${where}
|
||||
ORDER BY embeddings <-> $1 LIMIT ${max}
|
||||
`;
|
||||
|
||||
const db = (await this.getDb()) as pg.Client;
|
||||
const params = this.collection.length
|
||||
? [embedding, this.collection]
|
||||
: [embedding];
|
||||
const results = await db.query(sql, params);
|
||||
|
||||
const nodes = results.rows.map((row) => {
|
||||
return new Document({
|
||||
id_: row.id,
|
||||
text: row.document,
|
||||
metadata: row.metadata,
|
||||
embedding: row.embeddings,
|
||||
});
|
||||
});
|
||||
|
||||
const ret = {
|
||||
nodes: nodes,
|
||||
similarities: results.rows.map((row) => row.embeddings),
|
||||
ids: results.rows.map((row) => row.id),
|
||||
};
|
||||
|
||||
return Promise.resolve(ret);
|
||||
}
|
||||
|
||||
/**
|
||||
* Required by VectorStore interface. Currently ignored.
|
||||
* @param persistPath
|
||||
* @param fs
|
||||
* @returns Resolved Promise.
|
||||
*/
|
||||
persist(
|
||||
persistPath: string,
|
||||
fs?: GenericFileSystem | undefined,
|
||||
): Promise<void> {
|
||||
return Promise.resolve();
|
||||
}
|
||||
}
|
||||
@@ -1,11 +1,11 @@
|
||||
import _ from "lodash";
|
||||
import * as path from "path";
|
||||
import { BaseNode } from "../../Node";
|
||||
import {
|
||||
getTopKEmbeddings,
|
||||
getTopKEmbeddingsLearner,
|
||||
getTopKMMREmbeddings,
|
||||
} from "../../Embedding";
|
||||
import { BaseNode } from "../../Node";
|
||||
} from "../../embeddings";
|
||||
import { GenericFileSystem, exists } from "../FileSystem";
|
||||
import { DEFAULT_FS, DEFAULT_PERSIST_DIR } from "../constants";
|
||||
import {
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import { BaseNode } from "../../Node";
|
||||
import { GenericFileSystem } from "../FileSystem";
|
||||
|
||||
export interface VectorStoreQueryResult {
|
||||
nodes?: BaseNode[];
|
||||
@@ -62,7 +61,9 @@ export interface VectorStore {
|
||||
isEmbeddingQuery?: boolean;
|
||||
client(): any;
|
||||
add(embeddingResults: BaseNode[]): Promise<string[]>;
|
||||
delete(refDocId: string, deleteKwargs?: any): Promise<void>;
|
||||
query(query: VectorStoreQuery, kwargs?: any): Promise<VectorStoreQueryResult>;
|
||||
persist(persistPath: string, fs?: GenericFileSystem): Promise<void>;
|
||||
delete(refDocId: string, deleteOptions?: any): Promise<void>;
|
||||
query(
|
||||
query: VectorStoreQuery,
|
||||
options?: any,
|
||||
): Promise<VectorStoreQueryResult>;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
import { BaseNode, jsonToNode, Metadata, ObjectType } from "../../Node";
|
||||
|
||||
const DEFAULT_TEXT_KEY = "text";
|
||||
|
||||
export function validateIsFlat(obj: { [key: string]: any }): void {
|
||||
for (let key in obj) {
|
||||
if (typeof obj[key] === "object" && obj[key] !== null) {
|
||||
throw new Error(`Value for metadata ${key} must not be another object`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function nodeToMetadata(
|
||||
node: BaseNode,
|
||||
removeText: boolean = false,
|
||||
textField: string = DEFAULT_TEXT_KEY,
|
||||
flatMetadata: boolean = false,
|
||||
): Metadata {
|
||||
const nodeObj = node.toJSON();
|
||||
const metadata = node.metadata;
|
||||
|
||||
if (flatMetadata) {
|
||||
validateIsFlat(node.metadata);
|
||||
}
|
||||
|
||||
if (removeText) {
|
||||
nodeObj[textField] = "";
|
||||
}
|
||||
|
||||
nodeObj["embedding"] = null;
|
||||
|
||||
metadata["_node_content"] = JSON.stringify(nodeObj);
|
||||
metadata["_node_type"] = node.constructor.name.replace("_", ""); // remove leading underscore to be compatible with Python
|
||||
|
||||
metadata["document_id"] = node.sourceNode?.nodeId || "None";
|
||||
metadata["doc_id"] = node.sourceNode?.nodeId || "None";
|
||||
metadata["ref_doc_id"] = node.sourceNode?.nodeId || "None";
|
||||
|
||||
return metadata;
|
||||
}
|
||||
|
||||
export function metadataDictToNode(metadata: Metadata): BaseNode {
|
||||
const nodeContent = metadata["_node_content"];
|
||||
if (!nodeContent) {
|
||||
throw new Error("Node content not found in metadata.");
|
||||
}
|
||||
const nodeObj = JSON.parse(nodeContent);
|
||||
|
||||
// Note: we're using the name of the class stored in `_node_type`
|
||||
// and not the type attribute to reconstruct
|
||||
// the node. This way we're compatible with LlamaIndex Python
|
||||
const node_type = metadata["_node_type"];
|
||||
switch (node_type) {
|
||||
case "IndexNode":
|
||||
return jsonToNode(nodeObj, ObjectType.INDEX);
|
||||
default:
|
||||
return jsonToNode(nodeObj, ObjectType.TEXT);
|
||||
}
|
||||
}
|
||||
@@ -1,18 +1,18 @@
|
||||
import { OpenAIEmbedding } from "../Embedding";
|
||||
import {
|
||||
CallbackManager,
|
||||
RetrievalCallbackResponse,
|
||||
StreamCallbackResponse,
|
||||
} from "../callbacks/CallbackManager";
|
||||
import { OpenAIEmbedding } from "../embeddings";
|
||||
import { SummaryIndex } from "../indices/summary";
|
||||
import { VectorStoreIndex } from "../indices/vectorStore/VectorStoreIndex";
|
||||
import { OpenAI } from "../llm/LLM";
|
||||
import { Document } from "../Node";
|
||||
import {
|
||||
ResponseSynthesizer,
|
||||
SimpleResponseBuilder,
|
||||
} from "../ResponseSynthesizer";
|
||||
import { ServiceContext, serviceContextFromDefaults } from "../ServiceContext";
|
||||
import {
|
||||
CallbackManager,
|
||||
RetrievalCallbackResponse,
|
||||
StreamCallbackResponse,
|
||||
} from "../callbacks/CallbackManager";
|
||||
import { SummaryIndex } from "../indices/summary";
|
||||
import { VectorStoreIndex } from "../indices/vectorStore/VectorStoreIndex";
|
||||
import { OpenAI } from "../llm/LLM";
|
||||
import { mockEmbeddingModel, mockLlmGeneration } from "./utility/mockOpenAI";
|
||||
|
||||
// Mock the OpenAI getOpenAISession function during testing
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user