Compare commits

..

2 Commits

Author SHA1 Message Date
leehuwuj 4c4af6a0a8 e2e install pack and run the package 2024-02-06 10:20:46 +07:00
leehuwuj 4eecc5148e update create-llama e2e command 2024-02-06 09:45:50 +07:00
136 changed files with 1288 additions and 2706 deletions
-5
View File
@@ -1,5 +0,0 @@
---
"llamaindex": patch
---
feat: add filtering of metadata to PGVectorStore
-5
View File
@@ -1,5 +0,0 @@
---
"llamaindex": patch
---
feat(reranker): cohere reranker
+5
View File
@@ -0,0 +1,5 @@
---
"llamaindex": patch
---
fix: update `VectorIndexRetriever` constructor parameters' type.
-5
View File
@@ -1,5 +0,0 @@
---
"llamaindex": patch
---
feat: use batching in vector store index
-5
View File
@@ -1,5 +0,0 @@
---
"create-llama": patch
---
update fastapi for CVE-2024-24762
-5
View File
@@ -1,5 +0,0 @@
---
"llamaindex": patch
---
Add reader for LlamaParse
+2 -2
View File
@@ -4,6 +4,6 @@
"ghcr.io/devcontainers/features/node:1": {},
"ghcr.io/devcontainers-contrib/features/turborepo-npm:1": {},
"ghcr.io/devcontainers-contrib/features/typescript:2": {},
"ghcr.io/devcontainers-contrib/features/pnpm:2": {}
}
"ghcr.io/devcontainers-contrib/features/pnpm:2": {},
},
}
-2
View File
@@ -1,2 +0,0 @@
examples/readers/data/** binary
examples/data/** binary
+1 -7
View File
@@ -49,14 +49,8 @@ jobs:
- name: Build create-llama
run: pnpm run build
working-directory: ./packages/create-llama
- name: Pack
run: pnpm pack --pack-destination ./output
working-directory: ./packages/create-llama
- name: Extract Pack
run: tar -xvzf ./output/*.tgz -C ./output
working-directory: ./packages/create-llama
- name: Run Playwright tests
run: pnpm exec playwright test
run: pnpm run e2e
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
working-directory: ./packages/create-llama
-6
View File
@@ -38,12 +38,6 @@ jobs:
- name: Run Circular Dependency Check
run: pnpm run circular-check
working-directory: ./packages/core
- uses: actions/upload-artifact@v3
if: failure()
with:
name: typecheck-build-dist
path: ./packages/core/dist
if-no-files-found: error
typecheck-examples:
runs-on: ubuntu-latest
-12
View File
@@ -78,15 +78,3 @@ pnpm start
That should start a webserver which will serve the docs on https://localhost:3000
Any changes you make should be reflected in the browser. If you need to regenerate the API docs and find that your TSDoc isn't getting the updates, feel free to remove apps/docs/api. It will automatically regenerate itself when you run pnpm start again.
## Publishing
To publish a new version of the library, run
```shell
pnpm new-llamaindex
pnpm new-create-llama
pnpm release
git push # push to the main branch
git push --tags
```
@@ -1,128 +0,0 @@
# OpenAI Agent + QueryEngineTool
QueryEngineTool is a tool that allows you to query a vector index. In this example, we will create a vector index from a set of documents and then create a QueryEngineTool from the vector index. We will then create an OpenAIAgent with the QueryEngineTool and chat with the agent.
## Setup
First, you need to install the `llamaindex` package. You can do this by running the following command in your terminal:
```bash
pnpm i llamaindex
```
Then you can import the necessary classes and functions.
```ts
import {
OpenAIAgent,
SimpleDirectoryReader,
VectorStoreIndex,
QueryEngineTool,
} from "llamaindex";
```
## Create a vector index
Now we can create a vector index from a set of documents.
```ts
// Load the documents
const documents = await new SimpleDirectoryReader().loadData({
directoryPath: "node_modules/llamaindex/examples/",
});
// Create a vector index from the documents
const vectorIndex = await VectorStoreIndex.fromDocuments(documents);
```
## Create a QueryEngineTool
Now we can create a QueryEngineTool from the vector index.
```ts
// Create a query engine from the vector index
const abramovQueryEngine = vectorIndex.asQueryEngine();
// Create a QueryEngineTool with the query engine
const queryEngineTool = new QueryEngineTool({
queryEngine: abramovQueryEngine,
metadata: {
name: "abramov_query_engine",
description: "A query engine for the Abramov documents",
},
});
```
## Create an OpenAIAgent
```ts
// Create an OpenAIAgent with the query engine tool tools
const agent = new OpenAIAgent({
tools: [queryEngineTool],
verbose: true,
});
```
## Chat with the agent
Now we can chat with the agent.
```ts
const response = await agent.chat({
message: "What was his salary?",
});
console.log(String(response));
```
## Full code
```ts
import {
OpenAIAgent,
SimpleDirectoryReader,
VectorStoreIndex,
QueryEngineTool,
} from "llamaindex";
async function main() {
// Load the documents
const documents = await new SimpleDirectoryReader().loadData({
directoryPath: "node_modules/llamaindex/examples/",
});
// Create a vector index from the documents
const vectorIndex = await VectorStoreIndex.fromDocuments(documents);
// Create a query engine from the vector index
const abramovQueryEngine = vectorIndex.asQueryEngine();
// Create a QueryEngineTool with the query engine
const queryEngineTool = new QueryEngineTool({
queryEngine: abramovQueryEngine,
metadata: {
name: "abramov_query_engine",
description: "A query engine for the Abramov documents",
},
});
// Create an OpenAIAgent with the function tools
const agent = new OpenAIAgent({
tools: [queryEngineTool],
verbose: true,
});
// Chat with the agent
const response = await agent.chat({
message: "What was his salary?",
});
// Print the response
console.log(String(response));
}
main().then(() => {
console.log("Done");
});
```
+17
View File
@@ -0,0 +1,17 @@
---
sidebar_position: 3
---
# Reader / Loader
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";
documents = new SimpleDirectoryReader().loadData("./data");
```
## API Reference
- [SimpleDirectoryReader](../api/classes/SimpleDirectoryReader.md)
-35
View File
@@ -1,35 +0,0 @@
---
sidebar_position: 4
---
import CodeBlock from "@theme/CodeBlock";
import CodeSource from "!raw-loader!../../../../examples/readers/src/simple-directory-reader";
import CodeSource2 from "!raw-loader!../../../../examples/readers/src/custom-simple-directory-reader";
# Loader
Before you can start indexing your documents, you need to load them into memory.
### SimpleDirectoryReader
[![Open in StackBlitz](https://developer.stackblitz.com/img/open_in_stackblitz.svg)](https://stackblitz.com/github/run-llama/LlamaIndexTS/tree/main/examples/readers?file=src/simple-directory-reader.ts&title=Simple%20Directory%20Reader)
LlamaIndex.TS supports easy loading of files from folders using the `SimpleDirectoryReader` class.
It is a simple reader that reads all files from a directory and its subdirectories.
<CodeBlock language="ts">{CodeSource}</CodeBlock>
Currently, it supports reading `.csv`, `.docx`, `.html`, `.md` and `.pdf` files,
but support for other file types is planned.
Also, you can provide a `defaultReader` as a fallback for files with unsupported extensions.
Or pass new readers for `fileExtToReader` to support more file types.
<CodeBlock language="ts" showLineNumbers metastring="{8-12,17-21}">
{CodeSource2}
</CodeBlock>
## API Reference
- [SimpleDirectoryReader](../api/classes/SimpleDirectoryReader.md)
+1 -1
View File
@@ -1,5 +1,5 @@
---
sidebar_position: 4
sidebar_position: 3
---
# Embedding
@@ -2,7 +2,7 @@
sidebar_position: 3
---
# Large Language Models (LLMs)
# LLM
The LLM is responsible for reading text and generating natural language responses to queries. By default, LlamaIndex.TS uses `gpt-3.5-turbo`.
@@ -1,2 +0,0 @@
label: "LLMs"
position: 3
@@ -1 +0,0 @@
label: "Available LLMs"
@@ -1,80 +0,0 @@
# Anthropic
## Usage
```ts
import { Anthropic, serviceContextFromDefaults } from "llamaindex";
const anthropicLLM = new Anthropic({
apiKey: "<YOUR_API_KEY>",
});
const serviceContext = serviceContextFromDefaults({ llm: anthropicLLM });
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Query
```ts
const queryEngine = index.asQueryEngine();
const query = "What is the meaning of life?";
const results = await queryEngine.query({
query,
});
```
## Full Example
```ts
import {
Anthropic,
Document,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
async function main() {
// Create an instance of the Anthropic LLM
const anthropicLLM = new Anthropic({
apiKey: "<YOUR_API_KEY>",
});
// Create a service context
const serviceContext = serviceContextFromDefaults({ llm: anthropicLLM });
const document = new Document({ text: essay, id_: "essay" });
// Load and index documents
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
// Create a query engine
const queryEngine = index.asQueryEngine({
retriever,
});
const query = "What is the meaning of life?";
// Query
const response = await queryEngine.query({
query,
});
// Log the response
console.log(response.response);
}
```
@@ -1,88 +0,0 @@
# Azure OpenAI
To use Azure OpenAI, you only need to set a few environment variables together with the `OpenAI` class.
For example:
## Environment Variables
```
export AZURE_OPENAI_KEY="<YOUR KEY HERE>"
export AZURE_OPENAI_ENDPOINT="<YOUR ENDPOINT, see https://learn.microsoft.com/en-us/azure/ai-services/openai/quickstart?tabs=command-line%2Cpython&pivots=rest-api>"
export AZURE_OPENAI_DEPLOYMENT="gpt-4" # or some other deployment name
```
## Usage
```ts
import { OpenAI, serviceContextFromDefaults } from "llamaindex";
const azureOpenaiLLM = new OpenAI({ model: "gpt-4", temperature: 0 });
const serviceContext = serviceContextFromDefaults({ llm: azureOpenaiLLM });
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Query
```ts
const queryEngine = index.asQueryEngine();
const query = "What is the meaning of life?";
const results = await queryEngine.query({
query,
});
```
## Full Example
```ts
import {
Anthropic,
Document,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
async function main() {
// Create an instance of the LLM
const azureOpenaiLLM = new OpenAI({ model: "gpt-4", temperature: 0 });
// Create a service context
const serviceContext = serviceContextFromDefaults({ llm: azureOpenaiLLM });
const document = new Document({ text: essay, id_: "essay" });
// Load and index documents
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
// Create a query engine
const queryEngine = index.asQueryEngine({
retriever,
});
const query = "What is the meaning of life?";
// Query
const response = await queryEngine.query({
query,
});
// Log the response
console.log(response.response);
}
```
@@ -1,97 +0,0 @@
# LLama2
## Usage
```ts
import { Ollama, serviceContextFromDefaults } from "llamaindex";
const llama2LLM = new LlamaDeuce({ chatStrategy: DeuceChatStrategy.META });
const serviceContext = serviceContextFromDefaults({ llm: llama2LLM });
```
## Usage with Replication
```ts
import {
Ollama,
ReplicateSession,
serviceContextFromDefaults,
} from "llamaindex";
const replicateSession = new ReplicateSession({
replicateKey,
});
const llama2LLM = new LlamaDeuce({
chatStrategy: DeuceChatStrategy.META,
replicateSession,
});
const serviceContext = serviceContextFromDefaults({ llm: llama2LLM });
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Query
```ts
const queryEngine = index.asQueryEngine();
const query = "What is the meaning of life?";
const results = await queryEngine.query({
query,
});
```
## Full Example
```ts
import {
Anthropic,
Document,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
async function main() {
// Create an instance of the LLM
const llama2LLM = new LlamaDeuce({ chatStrategy: DeuceChatStrategy.META });
// Create a service context
const serviceContext = serviceContextFromDefaults({ llm: mistralLLM });
const document = new Document({ text: essay, id_: "essay" });
// Load and index documents
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
// Create a query engine
const queryEngine = index.asQueryEngine({
retriever,
});
const query = "What is the meaning of life?";
// Query
const response = await queryEngine.query({
query,
});
// Log the response
console.log(response.response);
}
```
@@ -1,79 +0,0 @@
# Mistral
## Usage
```ts
import { Ollama, serviceContextFromDefaults } from "llamaindex";
const mistralLLM = new MistralAI({
model: "mistral-tiny",
apiKey: "<YOUR_API_KEY>",
});
const serviceContext = serviceContextFromDefaults({ llm: mistralLLM });
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Query
```ts
const queryEngine = index.asQueryEngine();
const query = "What is the meaning of life?";
const results = await queryEngine.query({
query,
});
```
## Full Example
```ts
import {
Anthropic,
Document,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
async function main() {
// Create an instance of the LLM
const mistralLLM = new MistralAI({ model: "mistral-tiny" });
// Create a service context
const serviceContext = serviceContextFromDefaults({ llm: mistralLLM });
const document = new Document({ text: essay, id_: "essay" });
// Load and index documents
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
// Create a query engine
const queryEngine = index.asQueryEngine({
retriever,
});
const query = "What is the meaning of life?";
// Query
const response = await queryEngine.query({
query,
});
// Log the response
console.log(response.response);
}
```
@@ -1,76 +0,0 @@
# Ollama
## Usage
```ts
import { Ollama, serviceContextFromDefaults } from "llamaindex";
const ollamaLLM = new Ollama({ model: "llama2", temperature: 0.75 });
const serviceContext = serviceContextFromDefaults({ llm: ollamaLLM });
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Query
```ts
const queryEngine = index.asQueryEngine();
const query = "What is the meaning of life?";
const results = await queryEngine.query({
query,
});
```
## Full Example
```ts
import {
Anthropic,
Document,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
async function main() {
// Create an instance of the LLM
const ollamaLLM = new Ollama({ model: "llama2", temperature: 0.75 });
// Create a service context
const serviceContext = serviceContextFromDefaults({ llm: ollamaLLM });
const document = new Document({ text: essay, id_: "essay" });
// Load and index documents
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
// Create a query engine
const queryEngine = index.asQueryEngine({
retriever,
});
const query = "What is the meaning of life?";
// Query
const response = await queryEngine.query({
query,
});
// Log the response
console.log(response.response);
}
```
@@ -1,80 +0,0 @@
# OpenAI
```ts
import { OpenAI, serviceContextFromDefaults } from "llamaindex";
const openaiLLM = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0, apiKey: <YOUR_API_KEY> });
const serviceContext = serviceContextFromDefaults({ llm: openaiLLM });
```
You can setup the apiKey on the environment variables, like:
```bash
export OPENAI_API_KEY="<YOUR_API_KEY>"
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Query
```ts
const queryEngine = index.asQueryEngine();
const query = "What is the meaning of life?";
const results = await queryEngine.query({
query,
});
```
## Full Example
```ts
import {
Anthropic,
Document,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
async function main() {
// Create an instance of the LLM
const openaiLLM = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0 });
// Create a service context
const serviceContext = serviceContextFromDefaults({ llm: openaiLLM });
const document = new Document({ text: essay, id_: "essay" });
// Load and index documents
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
// Create a query engine
const queryEngine = index.asQueryEngine({
retriever,
});
const query = "What is the meaning of life?";
// Query
const response = await queryEngine.query({
query,
});
// Log the response
console.log(response.response);
}
```
@@ -1,80 +0,0 @@
# Portkey LLM
## Usage
```ts
import { Portkey, serviceContextFromDefaults } from "llamaindex";
const portkeyLLM = new Portkey({
apiKey: "<YOUR_API_KEY>",
});
const serviceContext = serviceContextFromDefaults({ llm: portkeyLLM });
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Query
```ts
const queryEngine = index.asQueryEngine();
const query = "What is the meaning of life?";
const results = await queryEngine.query({
query,
});
```
## Full Example
```ts
import {
Anthropic,
Document,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
async function main() {
// Create an instance of the LLM
const portkeyLLM = new Portkey({
apiKey: "<YOUR_API_KEY>",
});
// Create a service context
const serviceContext = serviceContextFromDefaults({ llm: portkeyLLM });
const document = new Document({ text: essay, id_: "essay" });
// Load and index documents
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
// Create a query engine
const queryEngine = index.asQueryEngine({
retriever,
});
const query = "What is the meaning of life?";
// Query
const response = await queryEngine.query({
query,
});
// Log the response
console.log(response.response);
}
```
@@ -1,80 +0,0 @@
# Together LLM
## Usage
```ts
import { TogetherLLM, serviceContextFromDefaults } from "llamaindex";
const togetherLLM = new TogetherLLM({
apiKey: "<YOUR_API_KEY>",
});
const serviceContext = serviceContextFromDefaults({ llm: togetherLLM });
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Query
```ts
const queryEngine = index.asQueryEngine();
const query = "What is the meaning of life?";
const results = await queryEngine.query({
query,
});
```
## Full Example
```ts
import {
Anthropic,
Document,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
async function main() {
// Create an instance of the LLM
const togetherLLM = new TogetherLLM({
apiKey: "<YOUR_API_KEY>",
});
// Create a service context
const serviceContext = serviceContextFromDefaults({ llm: togetherLLM });
const document = new Document({ text: essay, id_: "essay" });
// Load and index documents
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
// Create a query engine
const queryEngine = index.asQueryEngine({
retriever,
});
const query = "What is the meaning of life?";
// Query
const response = await queryEngine.query({
query,
});
// Log the response
console.log(response.response);
}
```
+1 -1
View File
@@ -1,5 +1,5 @@
---
sidebar_position: 4
sidebar_position: 3
---
# NodeParser
@@ -1,2 +0,0 @@
label: "Node Postprocessors"
position: 3
@@ -1,71 +0,0 @@
# Cohere Reranker
The Cohere Reranker is a postprocessor that uses the Cohere API to rerank the results of a search query.
## Setup
Firstly, you will need to install the `llamaindex` package.
```bash
pnpm install llamaindex
```
Now, you will need to sign up for an API key at [Cohere](https://cohere.ai/). Once you have your API key you can import the necessary modules and create a new instance of the `CohereRerank` class.
```ts
import {
CohereRerank,
Document,
OpenAI,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
```
## Load and index documents
For this example, we will use a single document. In a real-world scenario, you would have multiple documents to index.
```ts
const document = new Document({ text: essay, id_: "essay" });
const serviceContext = serviceContextFromDefaults({
llm: new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.1 }),
});
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
```
## Increase similarity topK to retrieve more results
The default value for `similarityTopK` is 2. This means that only the most similar document will be returned. To retrieve more results, you can increase the value of `similarityTopK`.
```ts
const retriever = index.asRetriever();
retriever.similarityTopK = 5;
```
## Create a new instance of the CohereRerank class
Then you can create a new instance of the `CohereRerank` class and pass in your API key and the number of results you want to return.
```ts
const nodePostprocessor = new CohereRerank({
apiKey: "<COHERE_API_KEY>",
topN: 4,
});
```
## Create a query engine with the retriever and node postprocessor
```ts
const queryEngine = index.asQueryEngine({
retriever,
nodePostprocessors: [nodePostprocessor],
});
// log the response
const response = await queryEngine.query("Where did the author grown up?");
```
@@ -1,110 +0,0 @@
# Node Postprocessors
## Concept
Node postprocessors are a set of modules that take a set of nodes, and apply some kind of transformation or filtering before returning them.
In LlamaIndex, node postprocessors are most commonly applied within a query engine, after the node retrieval step and before the response synthesis step.
LlamaIndex offers several node postprocessors for immediate use, while also providing a simple API for adding your own custom postprocessors.
## Usage Pattern
An example of using a node postprocessors is below:
```ts
import {
Node,
NodeWithScore,
SimilarityPostprocessor,
CohereRerank,
} from "llamaindex";
const nodes: NodeWithScore[] = [
{
node: new TextNode({ text: "hello world" }),
score: 0.8,
},
{
node: new TextNode({ text: "LlamaIndex is the best" }),
score: 0.6,
},
];
// similarity postprocessor: filter nodes below 0.75 similarity score
const processor = new SimilarityPostprocessor({
similarityCutoff: 0.7,
});
const filteredNodes = processor.postprocessNodes(nodes);
// cohere rerank: rerank nodes given query using trained model
const reranker = new CohereRerank({
apiKey: "<COHERE_API_KEY>",
topN: 2,
});
const rerankedNodes = await reranker.postprocessNodes(nodes, "<user_query>");
console.log(filteredNodes, rerankedNodes);
```
Now you can use the `filteredNodes` and `rerankedNodes` in your application.
## Using Node Postprocessors in LlamaIndex
Most commonly, node-postprocessors will be used in a query engine, where they are applied to the nodes returned from a retriever, and before the response synthesis step.
### Using Node Postprocessors in a Query Engine
```ts
import { Node, NodeWithScore, SimilarityPostprocessor, CohereRerank } from "llamaindex";
const nodes: NodeWithScore[] = [
{
node: new TextNode({ text: "hello world" }),
score: 0.8,
},
{
node: new TextNode({ text: "LlamaIndex is the best" }),
score: 0.6,
}
];
// cohere rerank: rerank nodes given query using trained model
const reranker = new CohereRerank({
apiKey: "<COHERE_API_KEY>,
topN: 2,
})
const document = new Document({ text: "essay", id_: "essay" });
const serviceContext = serviceContextFromDefaults({
llm: new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.1 }),
});
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
const queryEngine = index.asQueryEngine({
nodePostprocessors: [processor, reranker],
});
// all node post-processors will be applied during each query
const response = await queryEngine.query("<user_query>");
```
### Using with retrieved nodes
```ts
import { SimilarityPostprocessor } from "llamaindex";
nodes = await index.asRetriever().retrieve("test query str");
const processor = new SimilarityPostprocessor({
similarityCutoff: 0.7,
});
const filteredNodes = processor.postprocessNodes(nodes);
```
+2 -2
View File
@@ -6,6 +6,6 @@
"composite": true,
"incremental": true,
"outDir": "./lib",
"tsBuildInfoFile": "./lib/.tsbuildinfo"
}
"tsBuildInfoFile": "./lib/.tsbuildinfo",
},
}
-46
View File
@@ -1,46 +0,0 @@
import {
OpenAIAgent,
QueryEngineTool,
SimpleDirectoryReader,
VectorStoreIndex,
} from "llamaindex";
async function main() {
// Load the documents
const documents = await new SimpleDirectoryReader().loadData({
directoryPath: "node_modules/llamaindex/examples/",
});
// Create a vector index from the documents
const vectorIndex = await VectorStoreIndex.fromDocuments(documents);
// Create a query engine from the vector index
const abramovQueryEngine = vectorIndex.asQueryEngine();
// Create a QueryEngineTool with the query engine
const queryEngineTool = new QueryEngineTool({
queryEngine: abramovQueryEngine,
metadata: {
name: "abramov_query_engine",
description: "A query engine for the Abramov documents",
},
});
// Create an OpenAIAgent with the function tools
const agent = new OpenAIAgent({
tools: [queryEngineTool],
verbose: true,
});
// Chat with the agent
const response = await agent.chat({
message: "What was his salary?",
});
// Print the response
console.log(String(response));
}
main().then(() => {
console.log("Done");
});
+1 -3
View File
@@ -1,9 +1,7 @@
import { Anthropic } from "llamaindex";
(async () => {
const anthropic = new Anthropic({
apiKey: process.env.ANTHROPIC_API_KEY,
});
const anthropic = new Anthropic();
const result = await anthropic.chat({
messages: [
{ content: "You want to talk in rhymes.", role: "system" },
+2 -11
View File
@@ -14,27 +14,18 @@ Here are two sample scripts which work well with the sample data in the Astra Po
- `ASTRA_DB_APPLICATION_TOKEN`: The generated app token for your Astra database
- `ASTRA_DB_ENDPOINT`: The API endpoint for your Astra database
- `ASTRA_DB_NAMESPACE`: (Optional) The namespace where your collection is stored defaults to `default_keyspace`
- `OPENAI_API_KEY`: Your OpenAI key
2. `cd` Into the `examples` directory
3. run `npm i`
## Example load and query
Loads and queries a simple vectorstore with some documents about Astra DB
run `ts-node astradb/example`
## Movie Reviews Example
### Load the data
## Load the data
This sample loads the same dataset of movie reviews as the Astra Portal sample dataset. (Feel free to load the data in your the Astra Data Explorer to compare)
run `ts-node astradb/load`
### Use RAG to Query the data
## Use RAG to Query the data
Check out your data in the Astra Data Explorer and change the sample query as you see fit.
-58
View File
@@ -1,58 +0,0 @@
import {
AstraDBVectorStore,
Document,
storageContextFromDefaults,
VectorStoreIndex,
} from "llamaindex";
const collectionName = "test_collection";
async function main() {
try {
const docs = [
new Document({
text: "AstraDB is built on Apache Cassandra",
metadata: {
id: 123,
foo: "bar",
},
}),
new Document({
text: "AstraDB is a NoSQL DB",
metadata: {
id: 456,
foo: "baz",
},
}),
new Document({
text: "AstraDB supports vector search",
metadata: {
id: 789,
foo: "qux",
},
}),
];
const astraVS = new AstraDBVectorStore();
await astraVS.create(collectionName, {
vector: { dimension: 1536, metric: "cosine" },
});
await astraVS.connect(collectionName);
const ctx = await storageContextFromDefaults({ vectorStore: astraVS });
const index = await VectorStoreIndex.fromDocuments(docs, {
storageContext: ctx,
});
const queryEngine = index.asQueryEngine();
const response = await queryEngine.query({
query: "Describe AstraDB.",
});
console.log(response.toString());
} catch (e) {
console.error(e);
}
}
main();
+2 -2
View File
@@ -10,9 +10,9 @@ const collectionName = "movie_reviews";
async function main() {
try {
const reader = new PapaCSVReader(false);
const docs = await reader.loadData("./data/movie_reviews.csv");
const docs = await reader.loadData("../data/movie_reviews.csv");
const astraVS = new AstraDBVectorStore({ contentKey: "reviewtext" });
const astraVS = new AstraDBVectorStore();
await astraVS.create(collectionName, {
vector: { dimension: 1536, metric: "cosine" },
});
+2 -2
View File
@@ -8,7 +8,7 @@ const collectionName = "movie_reviews";
async function main() {
try {
const astraVS = new AstraDBVectorStore({ contentKey: "reviewtext" });
const astraVS = new AstraDBVectorStore();
await astraVS.connect(collectionName);
const ctx = serviceContextFromDefaults();
@@ -19,7 +19,7 @@ async function main() {
const queryEngine = await index.asQueryEngine({ retriever });
const results = await queryEngine.query({
query: 'How was "La Sapienza" reviewed?',
query: "What is the best reviewed movie?",
});
console.log(results.response);
Binary file not shown.
+59 -3
View File
@@ -1,5 +1,61 @@
## LlamaIndex Reader Examples
## Reader Examples
```shell
npm run start
These examples show how to use a specific reader class by loading a document and running a test query.
1. Make sure you are in `examples` directory
```bash
cd ./examples
```
2. Prepare `OPENAI_API_KEY` environment variable:
```bash
export OPENAI_API_KEY=your_openai_api_key
```
3. Run the following command to load documents and test query:
- MarkdownReader Example
```bash
npx ts-node readers/load-md.ts
```
- DocxReader Example
```bash
npx ts-node readers/load-docx.ts
```
- PdfReader Example
```bash
npx ts-node readers/load-pdf.ts
```
- HtmlReader Example
```bash
npx ts-node readers/load-html.ts
```
- CsvReader Example
```bash
npx ts-node readers/load-csv.ts
```
- NotionReader Example
```bash
export NOTION_TOKEN=your_notion_token
npx ts-node readers/load-notion.ts
```
- AssemblyAI Example
```bash
export ASSEMBLYAI_API_KEY=your_assemblyai_api_key
npx ts-node readers/load-assemblyai.ts
```
@@ -2,7 +2,7 @@ import { program } from "commander";
import { TranscribeParams, VectorStoreIndex } from "llamaindex";
import { AudioTranscriptReader } from "llamaindex/readers/AssemblyAIReader";
import { stdin as input, stdout as output } from "node:process";
import { createInterface } from "node:readline/promises";
import readline from "node:readline/promises";
program
.option("-a, --audio [string]", "URL or path of the audio file to transcribe")
@@ -35,7 +35,7 @@ program
// Create query engine
const queryEngine = index.asQueryEngine();
const rl = createInterface({ input, output });
const rl = readline.createInterface({ input, output });
while (true) {
const query = await rl.question("Ask a question: ");
@@ -10,7 +10,7 @@ import { PapaCSVReader } from "llamaindex/readers/CSVReader";
async function main() {
// Load CSV
const reader = new PapaCSVReader();
const path = "../data/titanic_train.csv";
const path = "data/titanic_train.csv";
const documents = await reader.loadData(path);
const serviceContext = serviceContextFromDefaults({
@@ -1,7 +1,7 @@
import { VectorStoreIndex } from "llamaindex";
import { DocxReader } from "llamaindex/readers/DocxReader";
const FILE_PATH = "../data/stars.docx";
const FILE_PATH = "./data/stars.docx";
const SAMPLE_QUERY = "Information about Zodiac";
async function main() {
@@ -4,7 +4,7 @@ import { HTMLReader } from "llamaindex/readers/HTMLReader";
async function main() {
// Load page
const reader = new HTMLReader();
const documents = await reader.loadData("../data/llamaindex.html");
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);
@@ -12,7 +12,7 @@ async function main() {
// Query the index
const queryEngine = index.asQueryEngine();
const response = await queryEngine.query({
query: "What can I do with LlamaIndex?",
query: "What were the notable changes in 18.1?",
});
// Output response
@@ -1,7 +1,7 @@
import { VectorStoreIndex } from "llamaindex";
import { MarkdownReader } from "llamaindex/readers/MarkdownReader";
const FILE_PATH = "../data/planets.md";
const FILE_PATH = "./data/planets.md";
const SAMPLE_QUERY = "List all planets";
async function main() {
@@ -3,7 +3,7 @@ import { program } from "commander";
import { VectorStoreIndex } from "llamaindex";
import { NotionReader } from "llamaindex/readers/NotionReader";
import { stdin as input, stdout as output } from "node:process";
import { createInterface } from "node:readline/promises";
import readline from "node:readline/promises";
program
.argument("[page]", "Notion page id (must be provided)")
@@ -70,7 +70,7 @@ program
// Create query engine
const queryEngine = index.asQueryEngine();
const rl = createInterface({ input, output });
const rl = readline.createInterface({ input, output });
while (true) {
const query = await rl.question("Query: ");
@@ -1,10 +1,13 @@
import { VectorStoreIndex } from "llamaindex";
import { PDFReader } from "llamaindex/readers/PDFReader";
import { resolve } from "node:path";
async function main() {
// Load PDF
const reader = new PDFReader();
const documents = await reader.loadData("../data/brk-2022.pdf");
const documents = await reader.loadData(
resolve(__dirname, "../data/brk-2022.pdf"),
);
// Split text and create embeddings. Store them in a VectorStoreIndex
const index = await VectorStoreIndex.fromDocuments(documents);
-22
View File
@@ -1,22 +0,0 @@
{
"name": "llamaindex-loader-example",
"private": true,
"type": "module",
"scripts": {
"start": "node --loader ts-node/esm ./src/simple-directory-reader.ts",
"start:csv": "node --loader ts-node/esm ./src/csv.ts",
"start:docx": "node --loader ts-node/esm ./src/docx.ts",
"start:html": "node --loader ts-node/esm ./src/html.ts",
"start:markdown": "node --loader ts-node/esm ./src/markdown.ts",
"start:pdf": "node --loader ts-node/esm ./src/pdf.ts",
"start:llamaparse": "node --loader ts-node/esm ./src/llamaparse.ts"
},
"dependencies": {
"llamaindex": "latest"
},
"devDependencies": {
"@types/node": "^20.11.14",
"ts-node": "^10.9.2",
"typescript": "^5.3.3"
}
}
@@ -1,26 +0,0 @@
import type { BaseReader, Document, Metadata } from "llamaindex";
import {
FILE_EXT_TO_READER,
SimpleDirectoryReader,
TextFileReader,
} from "llamaindex/readers/SimpleDirectoryReader";
class ZipReader implements BaseReader {
loadData(...args: any[]): Promise<Document<Metadata>[]> {
throw new Error("Implement me");
}
}
const reader = new SimpleDirectoryReader();
const documents = await reader.loadData({
directoryPath: "../data",
defaultReader: new TextFileReader(),
fileExtToReader: {
...FILE_EXT_TO_READER,
zip: new ZipReader(),
},
});
documents.forEach((doc) => {
console.log(`document (${doc.id_}):`, doc.getText());
});
-21
View File
@@ -1,21 +0,0 @@
import { LlamaParseReader, VectorStoreIndex } from "llamaindex";
async function main() {
// Load PDF using LlamaParse
const reader = new LlamaParseReader({ resultType: "markdown" });
const documents = await reader.loadData("../data/TOS.pdf");
// 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({
query: "What is the license grant in the TOS?",
});
// Output response
console.log(response.toString());
}
main().catch(console.error);
@@ -1,10 +0,0 @@
import { SimpleDirectoryReader } from "llamaindex/readers/SimpleDirectoryReader";
// or
// import { SimpleDirectoryReader } from 'llamaindex'
const reader = new SimpleDirectoryReader();
const documents = await reader.loadData("../data");
documents.forEach((doc) => {
console.log(`document (${doc.id_}):`, doc.getText());
});
-11
View File
@@ -1,11 +0,0 @@
{
"compilerOptions": {
"target": "es2017",
"module": "node16",
"moduleResolution": "node16",
"outDir": "./dist",
"types": ["node"],
"skipLibCheck": true
},
"include": ["./src/**/*.ts"]
}
-55
View File
@@ -1,55 +0,0 @@
import {
CohereRerank,
Document,
OpenAI,
VectorStoreIndex,
serviceContextFromDefaults,
} from "llamaindex";
import essay from "../essay";
async function main() {
const document = new Document({ text: essay, id_: "essay" });
const serviceContext = serviceContextFromDefaults({
llm: new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.1 }),
});
const index = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
const retriever = index.asRetriever();
retriever.similarityTopK = 5;
const nodePostprocessor = new CohereRerank({
apiKey: "<COHERE_API_KEY>",
topN: 5,
});
const queryEngine = index.asQueryEngine({
retriever,
nodePostprocessors: [nodePostprocessor],
});
const baseQueryEngine = index.asQueryEngine({
retriever,
});
const response = await queryEngine.query({
query: "What did the author do growing up?",
});
// cohere response
console.log(response.response);
const baseResponse = await baseQueryEngine.query({
query: "What did the author do growing up?",
});
// response without cohere
console.log(baseResponse.response);
}
main().catch(console.error);
+10 -17
View File
@@ -1,9 +1,4 @@
import {
Document,
QueryEngineTool,
SubQuestionQueryEngine,
VectorStoreIndex,
} from "llamaindex";
import { Document, SubQuestionQueryEngine, VectorStoreIndex } from "llamaindex";
import essay from "./essay";
@@ -11,18 +6,16 @@ import essay from "./essay";
const document = new Document({ text: essay, id_: essay });
const index = await VectorStoreIndex.fromDocuments([document]);
const queryEngineTools = [
new QueryEngineTool({
queryEngine: index.asQueryEngine(),
metadata: {
name: "pg_essay",
description: "Paul Graham essay on What I Worked On",
},
}),
];
const queryEngine = SubQuestionQueryEngine.fromDefaults({
queryEngineTools,
queryEngineTools: [
{
queryEngine: index.asQueryEngine(),
metadata: {
name: "pg_essay",
description: "Paul Graham essay on What I Worked On",
},
},
],
});
const response = await queryEngine.query({
+5 -5
View File
@@ -1,6 +1,6 @@
{
"compilerOptions": {
"target": "es2017",
"target": "es2016",
"module": "esnext",
"moduleResolution": "bundler",
"esModuleInterop": true,
@@ -10,13 +10,13 @@
"outDir": "./lib",
"tsBuildInfoFile": "./lib/.tsbuildinfo",
"incremental": true,
"composite": true
"composite": true,
},
"ts-node": {
"files": true,
"compilerOptions": {
"module": "commonjs"
}
"module": "commonjs",
},
},
"include": ["./**/*.ts"]
"include": ["./**/*.ts"],
}
+7 -6
View File
@@ -18,19 +18,20 @@
},
"devDependencies": {
"@changesets/cli": "^2.27.1",
"@types/jest": "^29.5.12",
"@turbo/gen": "^1.11.3",
"@types/jest": "^29.5.11",
"eslint": "^8.56.0",
"eslint-config-custom": "workspace:*",
"husky": "^9.0.10",
"husky": "^9.0.6",
"jest": "^29.7.0",
"lint-staged": "^15.2.2",
"prettier": "^3.2.5",
"lint-staged": "^15.2.0",
"prettier": "^3.2.4",
"prettier-plugin-organize-imports": "^3.2.4",
"ts-jest": "^29.1.2",
"turbo": "^1.12.3",
"turbo": "^1.11.3",
"typescript": "^5.3.3"
},
"packageManager": "pnpm@8.15.1",
"packageManager": "pnpm@8.14.3+sha256.2d0363bb6c314daa67087ef07743eea1ba2e2d360c835e8fec6b5575e4ed9484",
"pnpm": {
"overrides": {
"trim": "1.0.1",
-14
View File
@@ -1,19 +1,5 @@
# llamaindex
## 0.1.10
### Patch Changes
- b6c1500: feat(embedBatchSize): add batching for embeddings
- 6cc3a36: fix: update `VectorIndexRetriever` constructor parameters' type.
- cd82947: feat(queryEngineTool): add query engine tool to agents
## 0.1.9
### Patch Changes
- 09464e6: add OpenAIAgent (thanks @EmanuelCampos)
## 0.1.8
### Patch Changes
+19 -20
View File
@@ -1,27 +1,26 @@
{
"name": "llamaindex",
"private": true,
"version": "0.1.10",
"version": "0.1.8",
"license": "MIT",
"dependencies": {
"@anthropic-ai/sdk": "^0.13.0",
"@anthropic-ai/sdk": "^0.12.4",
"@datastax/astra-db-ts": "^0.1.4",
"@mistralai/mistralai": "^0.0.10",
"@notionhq/client": "^2.2.14",
"@pinecone-database/pinecone": "^1.1.3",
"@qdrant/js-client-rest": "^1.7.0",
"@xenova/transformers": "^2.15.0",
"assemblyai": "^4.2.2",
"@xenova/transformers": "^2.14.1",
"assemblyai": "^4.2.1",
"chromadb": "~1.7.3",
"cohere-ai": "^7.7.5",
"file-type": "^18.7.0",
"js-tiktoken": "^1.0.10",
"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.26.1",
"openai": "^4.26.0",
"papaparse": "^5.4.1",
"pathe": "^1.1.2",
"pdf2json": "^3.0.5",
@@ -30,18 +29,18 @@
"portkey-ai": "^0.1.16",
"rake-modified": "^1.0.8",
"replicate": "^0.25.2",
"string-strip-html": "^13.4.6",
"string-strip-html": "^13.4.5",
"wink-nlp": "^1.14.3"
},
"devDependencies": {
"@aws-crypto/sha256-js": "^5.2.0",
"@types/edit-json-file": "^1.7.3",
"@types/jest": "^29.5.12",
"@types/jest": "^29.5.11",
"@types/lodash": "^4.14.202",
"@types/node": "^18.19.14",
"@types/node": "^18.19.10",
"@types/papaparse": "^5.3.14",
"@types/pg": "^8.11.0",
"bunchee": "^4.4.6",
"bunchee": "^4.4.3",
"edit-json-file": "^1.8.0",
"madge": "^6.1.0",
"typescript": "^5.3.3"
@@ -119,6 +118,11 @@
"import": "./dist/Response.mjs",
"require": "./dist/Response.js"
},
"./Retriever": {
"types": "./dist/Retriever.d.mts",
"import": "./dist/Retriever.mjs",
"require": "./dist/Retriever.js"
},
"./ServiceContext": {
"types": "./dist/ServiceContext.d.mts",
"import": "./dist/ServiceContext.mjs",
@@ -129,15 +133,10 @@
"import": "./dist/TextSplitter.mjs",
"require": "./dist/TextSplitter.js"
},
"./tools": {
"types": "./dist/tools.d.mts",
"import": "./dist/tools.mjs",
"require": "./dist/tools.js"
},
"./readers": {
"types": "./dist/readers.d.mts",
"import": "./dist/readers.mjs",
"require": "./dist/readers.js"
"./Tool": {
"types": "./dist/Tool.d.mts",
"import": "./dist/Tool.mjs",
"require": "./dist/Tool.js"
},
"./readers/AssemblyAIReader": {
"types": "./dist/readers/AssemblyAIReader.d.mts",
+1 -8
View File
@@ -266,14 +266,7 @@ export class AgentRunner extends BaseAgentRunner {
let resultOutput;
while (true) {
const curStepOutput = await this._runStep(
task.taskId,
undefined,
ChatResponseMode.WAIT,
{
toolChoice,
},
);
const curStepOutput = await this._runStep(task.taskId);
if (curStepOutput.isLast) {
resultOutput = curStepOutput;
@@ -1,4 +1,3 @@
import type { Anthropic } from "@anthropic-ai/sdk";
import { NodeWithScore } from "../Node";
/*
@@ -40,7 +39,14 @@ export interface DefaultStreamToken {
//OpenAI stream token schema is the default.
//Note: Anthropic and Replicate also use similar token schemas.
export type OpenAIStreamToken = DefaultStreamToken;
export type AnthropicStreamToken = Anthropic.Completion;
export type AnthropicStreamToken = {
completion: string;
model: string;
stop_reason: string | undefined;
stop?: boolean | undefined;
log_id?: string;
};
//
//Callback Responses
//
@@ -36,7 +36,7 @@ export class HuggingFaceEmbedding extends BaseEmbedding {
return this.extractor;
}
override async getTextEmbedding(text: string): Promise<number[]> {
async getTextEmbedding(text: string): Promise<number[]> {
const extractor = await this.getExtractor();
const output = await extractor(text, { pooling: "mean", normalize: true });
return Array.from(output.data);
@@ -1,7 +0,0 @@
import { Ollama } from "../llm/ollama";
import { BaseEmbedding } from "./types";
/**
* OllamaEmbedding is an alias for Ollama that implements the BaseEmbedding interface.
*/
export class OllamaEmbedding extends Ollama implements BaseEmbedding {}
@@ -59,9 +59,7 @@ export class OpenAIEmbedding extends BaseEmbedding {
this.model = init?.model ?? "text-embedding-ada-002";
this.dimensions = init?.dimensions; // if no dimensions provided, will be undefined/not sent to OpenAI
this.embedBatchSize = init?.embedBatchSize ?? 10;
this.maxRetries = init?.maxRetries ?? 10;
this.timeout = init?.timeout ?? 60 * 1000; // Default is 60 seconds
this.additionalSessionOptions = init?.additionalSessionOptions;
@@ -102,43 +100,21 @@ export class OpenAIEmbedding extends BaseEmbedding {
}
}
/**
* Get embeddings for a batch of texts
* @param texts
* @param options
*/
private async getOpenAIEmbedding(input: string[]): Promise<number[][]> {
private async getOpenAIEmbedding(input: string) {
const { data } = await this.session.openai.embeddings.create({
model: this.model,
dimensions: this.dimensions, // only sent to OpenAI if set by user
input,
});
return data.map((d) => d.embedding);
return data[0].embedding;
}
/**
* Get embeddings for a batch of texts
* @param texts
*/
async getTextEmbeddings(texts: string[]): Promise<number[][]> {
return await this.getOpenAIEmbedding(texts);
}
/**
* Get embeddings for a single text
* @param texts
*/
async getTextEmbedding(text: string): Promise<number[]> {
return (await this.getOpenAIEmbedding([text]))[0];
return this.getOpenAIEmbedding(text);
}
/**
* Get embeddings for a query
* @param texts
* @param options
*/
async getQueryEmbedding(query: string): Promise<number[]> {
return (await this.getOpenAIEmbedding([query]))[0];
return this.getOpenAIEmbedding(query);
}
}
-1
View File
@@ -2,7 +2,6 @@ export * from "./ClipEmbedding";
export * from "./HuggingFaceEmbedding";
export * from "./MistralAIEmbedding";
export * from "./MultiModalEmbedding";
export { OllamaEmbedding } from "./OllamaEmbedding";
export * from "./OpenAIEmbedding";
export { TogetherEmbedding } from "./together";
export * from "./types";
+4 -62
View File
@@ -2,11 +2,7 @@ import { BaseNode, MetadataMode } from "../Node";
import { TransformComponent } from "../ingestion";
import { SimilarityType, similarity } from "./utils";
const DEFAULT_EMBED_BATCH_SIZE = 10;
export abstract class BaseEmbedding implements TransformComponent {
embedBatchSize = DEFAULT_EMBED_BATCH_SIZE;
similarity(
embedding1: number[],
embedding2: number[],
@@ -18,66 +14,12 @@ export abstract class BaseEmbedding implements TransformComponent {
abstract getTextEmbedding(text: string): Promise<number[]>;
abstract getQueryEmbedding(query: string): Promise<number[]>;
/**
* Optionally override this method to retrieve multiple embeddings in a single request
* @param texts
*/
async getTextEmbeddings(texts: string[]): Promise<Array<number[]>> {
const embeddings: number[][] = [];
for (const text of texts) {
const embedding = await this.getTextEmbedding(text);
embeddings.push(embedding);
}
return embeddings;
}
/**
* Get embeddings for a batch of texts
* @param texts
* @param options
*/
async getTextEmbeddingsBatch(
texts: string[],
options?: {
logProgress?: boolean;
},
): Promise<Array<number[]>> {
const resultEmbeddings: Array<number[]> = [];
const chunkSize = this.embedBatchSize;
const queue: string[] = texts;
const curBatch: string[] = [];
for (let i = 0; i < queue.length; i++) {
curBatch.push(queue[i]);
if (i == queue.length - 1 || curBatch.length == chunkSize) {
const embeddings = await this.getTextEmbeddings(curBatch);
resultEmbeddings.push(...embeddings);
if (options?.logProgress) {
console.log(`getting embedding progress: ${i} / ${queue.length}`);
}
curBatch.length = 0;
}
}
return resultEmbeddings;
}
async transform(nodes: BaseNode[], _options?: any): Promise<BaseNode[]> {
const texts = nodes.map((node) => node.getContent(MetadataMode.EMBED));
const embeddings = await this.getTextEmbeddingsBatch(texts);
for (let i = 0; i < nodes.length; i++) {
nodes[i].embedding = embeddings[i];
for (const node of nodes) {
node.embedding = await this.getTextEmbedding(
node.getContent(MetadataMode.EMBED),
);
}
return nodes;
}
}
@@ -22,17 +22,11 @@ export class DefaultContextGenerator implements ContextGenerator {
this.nodePostprocessors = init.nodePostprocessors || [];
}
private async applyNodePostprocessors(nodes: NodeWithScore[], query: string) {
let nodesWithScore = nodes;
for (const postprocessor of this.nodePostprocessors) {
nodesWithScore = await postprocessor.postprocessNodes(
nodesWithScore,
query,
);
}
return nodesWithScore;
private applyNodePostprocessors(nodes: NodeWithScore[]) {
return this.nodePostprocessors.reduce(
(nodes, nodePostprocessor) => nodePostprocessor.postprocessNodes(nodes),
nodes,
);
}
async generate(message: string, parentEvent?: Event): Promise<Context> {
@@ -48,10 +42,7 @@ export class DefaultContextGenerator implements ContextGenerator {
parentEvent,
);
const nodes = await this.applyNodePostprocessors(
sourceNodesWithScore,
message,
);
const nodes = this.applyNodePostprocessors(sourceNodesWithScore);
return {
message: {
@@ -36,17 +36,11 @@ export class RetrieverQueryEngine implements BaseQueryEngine {
this.nodePostprocessors = nodePostprocessors || [];
}
private async applyNodePostprocessors(nodes: NodeWithScore[], query: string) {
let nodesWithScore = nodes;
for (const postprocessor of this.nodePostprocessors) {
nodesWithScore = await postprocessor.postprocessNodes(
nodesWithScore,
query,
);
}
return nodesWithScore;
private applyNodePostprocessors(nodes: NodeWithScore[]) {
return this.nodePostprocessors.reduce(
(nodes, nodePostprocessor) => nodePostprocessor.postprocessNodes(nodes),
nodes,
);
}
private async retrieve(query: string, parentEvent: Event) {
@@ -56,7 +50,7 @@ export class RetrieverQueryEngine implements BaseQueryEngine {
this.preFilters,
);
return await this.applyNodePostprocessors(nodes, query);
return this.applyNodePostprocessors(nodes);
}
query(params: QueryEngineParamsStreaming): Promise<AsyncIterable<Response>>;
@@ -14,9 +14,9 @@ import {
} from "../../synthesizers";
import {
BaseQueryEngine,
BaseTool,
QueryEngineParamsNonStreaming,
QueryEngineParamsStreaming,
QueryEngineTool,
ToolMetadata,
} from "../../types";
import { BaseQuestionGenerator, SubQuestion } from "./types";
@@ -27,23 +27,28 @@ import { BaseQuestionGenerator, SubQuestion } from "./types";
export class SubQuestionQueryEngine implements BaseQueryEngine {
responseSynthesizer: BaseSynthesizer;
questionGen: BaseQuestionGenerator;
queryEngines: BaseTool[];
queryEngines: Record<string, BaseQueryEngine>;
metadatas: ToolMetadata[];
constructor(init: {
questionGen: BaseQuestionGenerator;
responseSynthesizer: BaseSynthesizer;
queryEngineTools: BaseTool[];
queryEngineTools: QueryEngineTool[];
}) {
this.questionGen = init.questionGen;
this.responseSynthesizer =
init.responseSynthesizer ?? new ResponseSynthesizer();
this.queryEngines = init.queryEngineTools;
this.queryEngines = init.queryEngineTools.reduce<
Record<string, BaseQueryEngine>
>((acc, tool) => {
acc[tool.metadata.name] = tool.queryEngine;
return acc;
}, {});
this.metadatas = init.queryEngineTools.map((tool) => tool.metadata);
}
static fromDefaults(init: {
queryEngineTools: BaseTool[];
queryEngineTools: QueryEngineTool[];
questionGen?: BaseQuestionGenerator;
responseSynthesizer?: BaseSynthesizer;
serviceContext?: ServiceContext;
@@ -117,24 +122,13 @@ export class SubQuestionQueryEngine implements BaseQueryEngine {
): Promise<NodeWithScore | null> {
try {
const question = subQ.subQuestion;
const queryEngine = this.queryEngines[subQ.toolName];
const queryEngine = this.queryEngines.find(
(tool) => tool.metadata.name === subQ.toolName,
);
if (!queryEngine) {
return null;
}
const responseText = await queryEngine?.call?.({
const response = await queryEngine.query({
query: question,
parentEvent,
});
if (!responseText) {
return null;
}
const responseText = response.response;
const nodeText = `Sub question: ${question}\nResponse: ${responseText}`;
const node = new TextNode({ text: nodeText });
return { node, score: 0 };
+10 -1
View File
@@ -21,7 +21,16 @@ export * from "./ingestion";
export * from "./llm";
export * from "./nodeParsers";
export * from "./postprocessors";
export * from "./readers";
export * from "./readers/AssemblyAIReader";
export * from "./readers/CSVReader";
export * from "./readers/DocxReader";
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 "./selectors";
export * from "./storage";
export * from "./synthesizers";
@@ -166,14 +166,20 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
nodes: BaseNode[],
options?: { logProgress?: boolean },
): Promise<BaseNode[]> {
const texts = nodes.map((node) => node.getContent(MetadataMode.EMBED));
const embeddings = await this.embedModel.getTextEmbeddingsBatch(texts, {
logProgress: options?.logProgress,
});
return nodes.map((node, i) => {
node.embedding = embeddings[i];
return node;
});
const nodesWithEmbeddings: BaseNode[] = [];
for (let i = 0; i < nodes.length; ++i) {
const node = nodes[i];
if (options?.logProgress) {
console.log(`Getting embedding for node ${i + 1}/${nodes.length}`);
}
node.embedding = await this.embedModel.getTextEmbedding(
node.getContent(MetadataMode.EMBED),
);
nodesWithEmbeddings.push(node);
}
return nodesWithEmbeddings;
}
/**
@@ -1,5 +1,5 @@
import { BaseNode, Document } from "../Node";
import { BaseReader } from "../readers/type";
import { BaseReader } from "../readers/base";
import { BaseDocumentStore, VectorStore } from "../storage";
import { IngestionCache, getTransformationHash } from "./IngestionCache";
import { DocStoreStrategy, createDocStoreStrategy } from "./strategies";
+1 -1
View File
@@ -1,5 +1,5 @@
import { CallbackManager, Event } from "../callbacks/CallbackManager";
import { BaseEmbedding } from "../embeddings/types";
import { BaseEmbedding } from "../embeddings";
import { ok } from "../env";
import {
ChatMessage,
@@ -8,7 +8,7 @@ export class MetadataReplacementPostProcessor implements BaseNodePostprocessor {
this.targetMetadataKey = targetMetadataKey;
}
async postprocessNodes(nodes: NodeWithScore[]): Promise<NodeWithScore[]> {
postprocessNodes(nodes: NodeWithScore[]): NodeWithScore[] {
for (let n of nodes) {
n.node.setContent(
n.node.metadata[this.targetMetadataKey] ??
@@ -8,7 +8,7 @@ export class SimilarityPostprocessor implements BaseNodePostprocessor {
this.similarityCutoff = options?.similarityCutoff;
}
async postprocessNodes(nodes: NodeWithScore[]) {
postprocessNodes(nodes: NodeWithScore[]) {
if (this.similarityCutoff === undefined) return nodes;
const cutoff = this.similarityCutoff || 0;
@@ -1,4 +1,3 @@
export * from "./MetadataReplacementPostProcessor";
export * from "./SimilarityPostprocessor";
export * from "./rerankers";
export * from "./types";
@@ -1,82 +0,0 @@
import { CohereClient } from "cohere-ai";
import { MetadataMode, NodeWithScore } from "../../Node";
import { BaseNodePostprocessor } from "../types";
type CohereRerankOptions = {
topN?: number;
model?: string;
apiKey: string | null;
};
export class CohereRerank implements BaseNodePostprocessor {
topN: number = 2;
model: string = "rerank-english-v2.0";
apiKey: string | null = null;
private client: CohereClient | null = null;
/**
* Constructor for CohereRerank.
* @param topN Number of nodes to return.
*/
constructor({
topN = 2,
model = "rerank-english-v2.0",
apiKey = null,
}: CohereRerankOptions) {
if (apiKey === null) {
throw new Error("CohereRerank requires an API key");
}
this.topN = topN;
this.model = model;
this.apiKey = apiKey;
this.client = new CohereClient({
token: this.apiKey,
});
}
/**
* Reranks the nodes using the Cohere API.
* @param nodes Array of nodes with scores.
* @param query Query string.
*/
async postprocessNodes(
nodes: NodeWithScore[],
query?: string,
): Promise<NodeWithScore[]> {
if (this.client === null) {
throw new Error("CohereRerank client is null");
}
if (nodes.length === 0) {
return [];
}
if (query === undefined) {
throw new Error("CohereRerank requires a query");
}
const results = await this.client.rerank({
query,
model: this.model,
topN: this.topN,
documents: nodes.map((n) => n.node.getContent(MetadataMode.ALL)),
});
const newNodes: NodeWithScore[] = [];
for (const result of results.results) {
const node = nodes[result.index];
newNodes.push({
node: node.node,
score: result.relevanceScore,
});
}
return newNodes;
}
}
@@ -1 +0,0 @@
export * from "./CohereRerank";
+1 -10
View File
@@ -1,14 +1,5 @@
import { NodeWithScore } from "../Node";
export interface BaseNodePostprocessor {
/**
* Send message along with the class's current chat history to the LLM.
* This version returns a promise for asynchronous operation.
* @param nodes Array of nodes with scores.
* @param query Optional query string.
*/
postprocessNodes(
nodes: NodeWithScore[],
query?: string,
): Promise<NodeWithScore[]>;
postprocessNodes: (nodes: NodeWithScore[]) => NodeWithScore[];
}
@@ -7,7 +7,7 @@ import {
TranscriptSentence,
} from "assemblyai";
import { Document } from "../Node";
import { BaseReader } from "./type";
import { BaseReader } from "./base";
type AssemblyAIOptions = Partial<BaseServiceParams>;
@@ -39,7 +39,7 @@ abstract class AssemblyAIReader implements BaseReader {
this.client = new AssemblyAI(options as BaseServiceParams);
}
abstract loadData(params: TranscribeParams | string): Promise<Document[]>;
abstract loadData(...args: any[]): Promise<Document[]>;
protected async transcribeOrGetTranscript(params: TranscribeParams | string) {
if (typeof params === "string") {
+2 -2
View File
@@ -2,14 +2,14 @@ import Papa, { ParseConfig } from "papaparse";
import { Document } from "../Node";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { FileReader } from "./type";
import { BaseReader } from "./base";
/**
* papaparse-based csv parser
* @class CSVReader
* @implements BaseReader
*/
export class PapaCSVReader implements FileReader {
export class PapaCSVReader implements BaseReader {
private concatRows: boolean;
private colJoiner: string;
private rowJoiner: string;
+2 -2
View File
@@ -2,9 +2,9 @@ import mammoth from "mammoth";
import { Document } from "../Node";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { FileReader } from "./type";
import { BaseReader } from "./base";
export class DocxReader implements FileReader {
export class DocxReader implements BaseReader {
/** DocxParser */
async loadData(
file: string,
+2 -2
View File
@@ -1,7 +1,7 @@
import { Document } from "../Node";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { FileReader } from "./type";
import { BaseReader } from "./base";
/**
* Extract the significant text from an arbitrary HTML document.
@@ -10,7 +10,7 @@ import { FileReader } from "./type";
* All other tags are removed, and the inner text is kept intact.
* Html entities (e.g., &amp;) are not decoded.
*/
export class HTMLReader implements FileReader {
export class HTMLReader implements BaseReader {
/**
* Public method for this reader.
* Required by BaseReader interface.
+2 -2
View File
@@ -1,12 +1,12 @@
import { Document, ImageDocument } from "../Node";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { FileReader } from "./type";
import { BaseReader } from "./base";
/**
* Reads the content of an image file into a Document object (which stores the image file as a Blob).
*/
export class ImageReader implements FileReader {
export class ImageReader implements BaseReader {
/**
* Public method for this reader.
* Required by BaseReader interface.
@@ -1,120 +0,0 @@
import { Document } from "../Node";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { FileReader } from "./type";
type ResultType = "text" | "markdown";
/**
* Represents a reader for parsing files using the LlamaParse API.
* See https://github.com/run-llama/llama_parse
*/
export class LlamaParseReader implements FileReader {
// The API key for the LlamaParse API.
apiKey: string;
// The base URL of the Llama Parsing API.
baseUrl: string = "https://api.cloud.llamaindex.ai/api/parsing";
// The maximum timeout in seconds to wait for the parsing to finish.
maxTimeout = 2000;
// The interval in seconds to check if the parsing is done.
checkInterval = 1;
// Whether to print the progress of the parsing.
verbose = true;
resultType: ResultType = "text";
constructor(params: Partial<LlamaParseReader> = {}) {
Object.assign(this, params);
params.apiKey = params.apiKey ?? process.env.LLAMA_CLOUD_API_KEY;
if (!params.apiKey) {
throw new Error(
"API Key is required for LlamaParseReader. Please pass the apiKey parameter or set the LLAMA_CLOUD_API_KEY environment variable.",
);
}
this.apiKey = params.apiKey;
}
async loadData(
file: string,
fs: GenericFileSystem = defaultFS,
): Promise<Document[]> {
if (!file.endsWith(".pdf")) {
throw new Error("Currently, only PDF files are supported.");
}
const metadata = { file_path: file };
// Load data, set the mime type
const data = await fs.readRawFile(file);
const mimeType = await this.getMimeType(data);
const body = new FormData();
body.set("file", new Blob([data], { type: mimeType }), file);
const headers = {
Authorization: `Bearer ${this.apiKey}`,
};
// Send the request, start job
const url = `${this.baseUrl}/upload`;
let response = await fetch(url, {
signal: AbortSignal.timeout(this.maxTimeout * 1000),
method: "POST",
body,
headers,
});
if (!response.ok) {
throw new Error(`Failed to parse the PDF file: ${await response.text()}`);
}
const jsonResponse = await response.json();
// Check the status of the job, return when done
const jobId = jsonResponse.id;
if (this.verbose) {
console.log(`Started parsing the file under job id ${jobId}`);
}
const resultUrl = `${this.baseUrl}/job/${jobId}/result/${this.resultType}`;
let start = Date.now();
let tries = 0;
while (true) {
await new Promise((resolve) =>
setTimeout(resolve, this.checkInterval * 1000),
);
response = await fetch(resultUrl, {
headers,
signal: AbortSignal.timeout(this.maxTimeout * 1000),
});
if (!response.ok) {
const end = Date.now();
if (end - start > this.maxTimeout * 1000) {
throw new Error(
`Timeout while parsing the PDF file: ${await response.text()}`,
);
}
if (this.verbose && tries % 10 === 0) {
process.stdout.write(".");
}
tries++;
continue;
}
const resultJson = await response.json();
return [
new Document({
text: resultJson[this.resultType],
metadata: metadata,
}),
];
}
}
private async getMimeType(data: Buffer): Promise<string> {
const { fileTypeFromBuffer } = await import("file-type");
const type = await fileTypeFromBuffer(data);
if (type?.mime !== "application/pdf") {
throw new Error("Currently, only PDF files are supported.");
}
return type.mime;
}
}
+2 -2
View File
@@ -1,7 +1,7 @@
import { Document } from "../Node";
import { defaultFS } from "../env";
import { GenericFileSystem } from "../storage";
import { FileReader } from "./type";
import { BaseReader } from "./base";
type MarkdownTuple = [string | null, string];
@@ -9,7 +9,7 @@ type MarkdownTuple = [string | null, string];
* Extract text from markdown files.
* Returns dictionary with keys as headers and values as the text between headers.
*/
export class MarkdownReader implements FileReader {
export class MarkdownReader implements BaseReader {
private _removeHyperlinks: boolean;
private _removeImages: boolean;
+2 -6
View File
@@ -1,7 +1,7 @@
import { Client } from "@notionhq/client";
import { crawler, Crawler, Pages, pageToString } from "notion-md-crawler";
import { Document } from "../Node";
import { BaseReader } from "./type";
import { BaseReader } from "./base";
type OptionalSerializers = Parameters<Crawler>[number]["serializers"];
@@ -42,11 +42,7 @@ export class NotionReader implements BaseReader {
toDocuments(pages: Pages): Document[] {
return Object.values(pages).map((page) => {
const text = pageToString(page);
return new Document({
id_: page.metadata.id, // Use the Notion-provided UUID for the document
text,
metadata: page.metadata,
});
return new Document({ text, metadata: page.metadata });
});
}
+1 -1
View File
@@ -1,7 +1,7 @@
import { Document } from "../Node";
import { createSHA256, defaultFS } from "../env";
import { GenericFileSystem } from "../storage/FileSystem";
import { BaseReader } from "./type";
import { BaseReader } from "./base";
/**
* Read the text of a PDF
@@ -1,5 +1,6 @@
import _ from "lodash";
import { Document } from "../Node";
import { defaultFS, path } from "../env";
import { defaultFS } from "../env";
import { CompleteFileSystem, walk } from "../storage/FileSystem";
import { PapaCSVReader } from "./CSVReader";
import { DocxReader } from "./DocxReader";
@@ -7,7 +8,7 @@ import { HTMLReader } from "./HTMLReader";
import { ImageReader } from "./ImageReader";
import { MarkdownReader } from "./MarkdownReader";
import { PDFReader } from "./PDFReader";
import { BaseReader } from "./type";
import { BaseReader } from "./base";
type ReaderCallback = (
category: "file" | "directory",
@@ -56,17 +57,13 @@ export type SimpleDirectoryReaderLoadDataParams = {
};
/**
* Read all the documents in a directory.
* Read all of the documents in a directory.
* By default, supports the list of file types
* in the FILE_EXT_TO_READER map.
*/
export class SimpleDirectoryReader implements BaseReader {
constructor(private observer?: ReaderCallback) {}
async loadData(
params: SimpleDirectoryReaderLoadDataParams,
): Promise<Document[]>;
async loadData(directoryPath: string): Promise<Document[]>;
async loadData(
params: SimpleDirectoryReaderLoadDataParams | string,
): Promise<Document[]> {
@@ -91,7 +88,7 @@ export class SimpleDirectoryReader implements BaseReader {
let docs: Document[] = [];
for await (const filePath of walk(fs, directoryPath)) {
try {
const fileExt = path.extname(filePath).slice(1).toLowerCase();
const fileExt = _.last(filePath.split(".")) || "";
// Observer can decide to skip each file
if (!this.doObserverCheck("file", filePath, ReaderStatus.STARTED)) {
@@ -99,11 +96,11 @@ export class SimpleDirectoryReader implements BaseReader {
continue;
}
let reader: BaseReader;
let reader = null;
if (fileExt in fileExtToReader) {
reader = fileExtToReader[fileExt];
} else if (defaultReader != null) {
} else if (!_.isNil(defaultReader)) {
reader = defaultReader;
} else {
const msg = `No reader for file extension of ${filePath}`;
@@ -1,6 +1,6 @@
import { MongoClient } from "mongodb";
import { Document, Metadata } from "../Node";
import { BaseReader } from "./type";
import { BaseReader } from "./base";
/**
* Read in from MongoDB
+8
View File
@@ -0,0 +1,8 @@
import { Document } from "../Node";
/**
* A reader takes imports data into Document objects.
*/
export interface BaseReader {
loadData(...args: any[]): Promise<Document[]>;
}
-12
View File
@@ -1,12 +0,0 @@
export * from "./AssemblyAIReader";
export * from "./CSVReader";
export * from "./DocxReader";
export * from "./HTMLReader";
export * from "./ImageReader";
export * from "./LlamaParseReader";
export * from "./MarkdownReader";
export * from "./NotionReader";
export * from "./PDFReader";
export * from "./SimpleDirectoryReader";
export * from "./SimpleMongoReader";
export * from "./type";
-16
View File
@@ -1,16 +0,0 @@
import { Document } from "../Node";
import { CompleteFileSystem } from "../storage";
/**
* A reader takes imports data into Document objects.
*/
export interface BaseReader {
loadData(...args: unknown[]): Promise<Document[]>;
}
/**
* A reader takes file paths and imports data into Document objects.
*/
export interface FileReader extends BaseReader {
loadData(filePath: string, fs?: CompleteFileSystem): Promise<Document[]>;
}
@@ -1,9 +1,8 @@
import { AstraDB } from "@datastax/astra-db-ts";
import { Collection } from "@datastax/astra-db-ts/dist/collections";
import { CreateCollectionOptions } from "@datastax/astra-db-ts/dist/collections/options";
import { BaseNode, MetadataMode } from "../../Node";
import { BaseNode, Document, MetadataMode } from "../../Node";
import { VectorStore, VectorStoreQuery, VectorStoreQueryResult } from "./types";
import { metadataDictToNode, nodeToMetadata } from "./utils";
const MAX_INSERT_BATCH_SIZE = 20;
@@ -13,7 +12,7 @@ export class AstraDBVectorStore implements VectorStore {
astraDBClient: AstraDB;
idKey: string;
contentKey: string;
contentKey: string | undefined; // if undefined the entirety of the node aside from the id and embedding will be stored as content
metadataKey: string;
private collection: Collection | undefined;
@@ -23,7 +22,6 @@ export class AstraDBVectorStore implements VectorStore {
params?: {
token: string;
endpoint: string;
namespace: string;
};
},
) {
@@ -42,15 +40,11 @@ export class AstraDBVectorStore implements VectorStore {
if (!endpoint) {
throw new Error("Must specify ASTRA_DB_ENDPOINT via env variable.");
}
const namespace =
init?.params?.namespace ??
process.env.ASTRA_DB_NAMESPACE ??
"default_keyspace";
this.astraDBClient = new AstraDB(token, endpoint, namespace);
this.astraDBClient = new AstraDB(token, endpoint);
}
this.idKey = init?.idKey ?? "_id";
this.contentKey = init?.contentKey ?? "content";
this.contentKey = init?.contentKey;
this.metadataKey = init?.metadataKey ?? "metadata";
}
@@ -108,20 +102,12 @@ export class AstraDBVectorStore implements VectorStore {
if (!nodes || nodes.length === 0) {
return [];
}
const dataToInsert = nodes.map((node) => {
const metadata = nodeToMetadata(
node,
true,
this.contentKey,
this.flatMetadata,
);
return {
_id: node.id_,
$vector: node.getEmbedding(),
[this.idKey]: node.id_,
[this.contentKey]: node.getContent(MetadataMode.NONE),
[this.metadataKey]: metadata,
content: node.getContent(MetadataMode.ALL),
metadata: node.metadata,
};
});
@@ -136,10 +122,11 @@ export class AstraDBVectorStore implements VectorStore {
for (const batch of batchData) {
console.debug(`Inserting batch of size ${batch.length}`);
await collection.insertMany(batch);
const result = await collection.insertMany(batch);
}
return dataToInsert.map((node) => node?.[this.idKey] as string);
return dataToInsert.map((node) => node._id);
}
/**
@@ -198,24 +185,27 @@ export class AstraDBVectorStore implements VectorStore {
const similarities: number[] = [];
await cursor.forEach(async (row: Record<string, any>) => {
const {
$vector: embedding,
$similarity: similarity,
[this.idKey]: id,
[this.contentKey]: content,
[this.metadataKey]: metadata = {},
...rest
} = row;
const id = row[this.idKey];
const embedding = row.$vector;
const similarity = row.$similarity;
const metadata = row[this.metadataKey];
const node = metadataDictToNode(metadata, {
fallback: {
id,
text: content,
metadata,
...rest,
},
// Remove fields from content
delete row[this.idKey];
delete row.$similarity;
delete row.$vector;
delete row[this.metadataKey];
const content = this.contentKey
? row[this.contentKey]
: JSON.stringify(row);
const node = new Document({
id_: id,
text: content,
metadata: metadata ?? {},
embedding: embedding,
});
node.setContent(content);
ids.push(id);
similarities.push(similarity);
@@ -107,17 +107,17 @@ export class PGVectorStore implements VectorStore {
await db.query(`CREATE SCHEMA IF NOT EXISTS ${this.schemaName}`);
const tbl = `CREATE TABLE IF NOT EXISTS ${this.schemaName}.${this.tableName}(
id uuid DEFAULT gen_random_uuid() PRIMARY KEY,
id uuid DEFAULT gen_random_uuid() PRIMARY KEY,
external_id VARCHAR,
collection VARCHAR,
document TEXT,
metadata JSONB DEFAULT '{}',
embeddings VECTOR(${this.dimensions})
)`;
)`;
await db.query(tbl);
const idxs = `CREATE INDEX IF NOT EXISTS idx_${this.tableName}_external_id ON ${this.schemaName}.${this.tableName} (external_id);
CREATE INDEX IF NOT EXISTS idx_${this.tableName}_collection ON ${this.schemaName}.${this.tableName} (collection);`;
CREATE INDEX IF NOT EXISTS idx_${this.tableName}_collection ON ${this.schemaName}.${this.tableName} (collection);`;
await db.query(idxs);
// TODO add IVFFlat or HNSW indexing?
@@ -140,8 +140,8 @@ export class PGVectorStore implements VectorStore {
* @returns The result of the delete query.
*/
async clearCollection() {
const sql: string = `DELETE FROM ${this.schemaName}.${this.tableName}
WHERE collection = $1`;
const sql: string = `DELETE FROM ${this.schemaName}.${this.tableName}
WHERE collection = $1`;
const db = (await this.getDb()) as pg.Client;
const ret = await db.query(sql, [this.collection]);
@@ -184,9 +184,9 @@ export class PGVectorStore implements VectorStore {
return Promise.resolve([]);
}
const sql: string = `INSERT INTO ${this.schemaName}.${this.tableName}
(id, external_id, collection, document, metadata, embeddings)
VALUES ($1, $2, $3, $4, $5, $6)`;
const sql: string = `INSERT INTO ${this.schemaName}.${this.tableName}
(id, external_id, collection, document, metadata, embeddings)
VALUES ($1, $2, $3, $4, $5, $6)`;
const db = (await this.getDb()) as pg.Client;
const data = this.getDataToInsert(embeddingResults);
@@ -220,8 +220,8 @@ export class PGVectorStore implements VectorStore {
const collectionCriteria = this.collection.length
? "AND collection = $2"
: "";
const sql: string = `DELETE FROM ${this.schemaName}.${this.tableName}
WHERE id = $1 ${collectionCriteria}`;
const sql: string = `DELETE FROM ${this.schemaName}.${this.tableName}
WHERE id = $1 ${collectionCriteria}`;
const db = (await this.getDb()) as pg.Client;
const params = this.collection.length
@@ -248,21 +248,8 @@ export class PGVectorStore implements VectorStore {
const embedding = "[" + query.queryEmbedding?.join(",") + "]";
const max = query.similarityTopK ?? 2;
const whereClauses = this.collection.length ? ["collection = $2"] : [];
const params: Array<string | number> = this.collection.length
? [embedding, this.collection]
: [embedding];
query.filters?.filters.forEach((filter, index) => {
const paramIndex = params.length + 1;
whereClauses.push(`metadata->>'${filter.key}' = $${paramIndex}`);
params.push(filter.value);
});
const where =
whereClauses.length > 0 ? `WHERE ${whereClauses.join(" AND ")}` : "";
const where = this.collection.length ? "WHERE collection = $2" : "";
// TODO Add collection filter if set
const sql = `SELECT
v.*,
embeddings <-> $1 s
@@ -273,6 +260,9 @@ export class PGVectorStore implements VectorStore {
`;
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) => {
+4 -19
View File
@@ -36,16 +36,7 @@ export function nodeToMetadata(
return metadata;
}
type MetadataDictToNodeOptions = {
// If the metadata doesn't contain node content, use this object as a fallback, for usage see
// AstraDBVectorStore.ts
fallback: Record<string, any>;
};
export function metadataDictToNode(
metadata: Metadata,
options?: MetadataDictToNodeOptions,
): BaseNode {
export function metadataDictToNode(metadata: Metadata): BaseNode {
const {
_node_content: nodeContent,
_node_type: nodeType,
@@ -54,17 +45,11 @@ export function metadataDictToNode(
ref_doc_id,
...rest
} = metadata;
let nodeObj;
if (!nodeContent) {
if (options?.fallback) {
nodeObj = options?.fallback;
} else {
throw new Error("Node content not found in metadata.");
}
} else {
nodeObj = JSON.parse(nodeContent);
nodeObj.metadata = rest;
throw new Error("Node content not found in metadata.");
}
const nodeObj = JSON.parse(nodeContent);
nodeObj.metadata = rest;
// Note: we're using the name of the class stored in `_node_type`
// and not the type attribute to reconstruct
+1 -38
View File
@@ -1,12 +1,4 @@
import { OpenAIEmbedding, similarity, SimilarityType } from "../embeddings";
import { mockEmbeddingModel } from "./utility/mockOpenAI";
// Mock the OpenAI getOpenAISession function during testing
jest.mock("../llm/open_ai", () => {
return {
getOpenAISession: jest.fn().mockImplementation(() => null),
};
});
import { similarity, SimilarityType } from "../embeddings";
describe("similarity", () => {
test("throws error on mismatched lengths", () => {
@@ -50,32 +42,3 @@ describe("similarity", () => {
);
});
});
describe("[OpenAIEmbedding]", () => {
let embedModel: OpenAIEmbedding;
beforeAll(() => {
let openAIEmbedding = new OpenAIEmbedding();
mockEmbeddingModel(openAIEmbedding);
embedModel = openAIEmbedding;
});
test("getTextEmbedding", async () => {
const embedding = await embedModel.getTextEmbedding("hello");
expect(embedding.length).toEqual(6);
});
test("getTextEmbeddings", async () => {
const texts = ["hello", "world"];
const embeddings = await embedModel.getTextEmbeddings(texts);
expect(embeddings.length).toEqual(1);
});
test("getTextEmbeddingsBatch", async () => {
const texts = ["hello", "world"];
const embeddings = await embedModel.getTextEmbeddingsBatch(texts);
expect(embeddings.length).toEqual(1);
});
});
@@ -18,15 +18,15 @@ describe("MetadataReplacementPostProcessor", () => {
];
});
test("Replaces the content of each node with specified metadata key if it exists", async () => {
test("Replaces the content of each node with specified metadata key if it exists", () => {
nodes[0].node.metadata = { targetKey: "NewContent" };
const newNodes = await postProcessor.postprocessNodes(nodes);
const newNodes = postProcessor.postprocessNodes(nodes);
// Check if node content was replaced correctly
expect(newNodes[0].node.getContent(MetadataMode.NONE)).toBe("NewContent");
});
test("Retains the original content of each node if no metadata key is found", async () => {
const newNodes = await postProcessor.postprocessNodes(nodes);
test("Retains the original content of each node if no metadata key is found", () => {
const newNodes = postProcessor.postprocessNodes(nodes);
// Check if node content remained unchanged
expect(newNodes[0].node.getContent(MetadataMode.NONE)).toBe("OldContent");
});
@@ -90,11 +90,6 @@ export function mockEmbeddingModel(embedModel: OpenAIEmbedding) {
resolve([1, 0, 0, 0, 0, 0]);
});
});
jest.spyOn(embedModel, "getTextEmbeddings").mockImplementation(async (x) => {
return new Promise((resolve) => {
resolve([[1, 0, 0, 0, 0, 0]]);
});
});
jest.spyOn(embedModel, "getQueryEmbedding").mockImplementation(async (x) => {
return new Promise((resolve) => {
resolve([0, 1, 0, 0, 0, 0]);
@@ -1,54 +0,0 @@
import { BaseQueryEngine, BaseTool, ToolMetadata } from "../types";
export type QueryEngineToolParams = {
queryEngine: BaseQueryEngine;
metadata: ToolMetadata;
};
type QueryEngineCallParams = {
query: string;
};
const DEFAULT_NAME = "query_engine_tool";
const DEFAULT_DESCRIPTION =
"Useful for running a natural language query against a knowledge base and get back a natural language response.";
const DEFAULT_PARAMETERS = {
type: "object",
properties: {
query: {
type: "string",
description: "The query to search for",
},
},
required: ["query"],
};
export class QueryEngineTool implements BaseTool {
private queryEngine: BaseQueryEngine;
metadata: ToolMetadata;
constructor({ queryEngine, metadata }: QueryEngineToolParams) {
this.queryEngine = queryEngine;
this.metadata = {
name: metadata?.name ?? DEFAULT_NAME,
description: metadata?.description ?? DEFAULT_DESCRIPTION,
parameters: metadata?.parameters ?? DEFAULT_PARAMETERS,
};
}
async call(...args: QueryEngineCallParams[]): Promise<any> {
let queryStr: string;
if (args && args.length > 0) {
queryStr = String(args[0].query);
} else {
throw new Error(
"Cannot call query engine without specifying `input` parameter.",
);
}
const response = await this.queryEngine.query({ query: queryStr });
return response.response;
}
}
-1
View File
@@ -1,3 +1,2 @@
export * from "./QueryEngineTool";
export * from "./functionTool";
export * from "./types";
+7
View File
@@ -40,6 +40,13 @@ export interface BaseTool {
metadata: ToolMetadata;
}
/**
* A Tool that uses a QueryEngine.
*/
export interface QueryEngineTool extends BaseTool {
queryEngine: BaseQueryEngine;
}
/**
* An OutputParser is used to extract structured data from the raw output of the LLM.
*/

Some files were not shown because too many files have changed in this diff Show More