Compare commits

..

4 Commits

Author SHA1 Message Date
Yi Ding 259fe63ceb strong types for prompts 2023-08-29 12:33:46 -07:00
Yi Ding d1aa3b7982 more changes for the summary index 2023-08-29 09:41:23 -07:00
Yi Ding e4af7b3a53 renamed listindex to summaryindex
Better aligns with what the index is used for
2023-08-29 09:32:04 -07:00
Yi Ding 51bd392fed 0.0.23 2023-08-25 12:28:06 -07:00
26 changed files with 200 additions and 157 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"llamaindex": patch
---
Renamed ListIndex to SummaryIndex to better indicate its use.
+5
View File
@@ -0,0 +1,5 @@
---
"llamaindex": patch
---
Strong types for prompts.
-5
View File
@@ -1,5 +0,0 @@
---
"llamaindex": patch
---
Added Markdown Reader (huge shoutout to @swk777)
+8
View File
@@ -1,5 +1,13 @@
# simple
## 0.0.21
### Patch Changes
- Updated dependencies
- Updated dependencies [9d6b2ed]
- llamaindex@0.0.23
## 0.0.20
### Patch Changes
+1 -4
View File
@@ -4,7 +4,6 @@ import {
PapaCSVReader,
ResponseSynthesizer,
serviceContextFromDefaults,
SimplePrompt,
VectorStoreIndex,
} from "llamaindex";
@@ -23,9 +22,7 @@ async function main() {
serviceContext,
});
const csvPrompt: SimplePrompt = (input) => {
const { context = "", query = "" } = input;
const csvPrompt = ({ context = "", query = "" }) => {
return `The following CSV file is loaded from ${path}
\`\`\`csv
${context}
+1 -8
View File
@@ -1,14 +1,7 @@
import { OpenAI } from "llamaindex";
(async () => {
const llm = new OpenAI({
model: "gpt-3.5-turbo",
temperature: 0.1,
additionalChatOptions: { frequency_penalty: 0.1 },
additionalSessionOptions: {
defaultHeaders: { "X-Test-Header-Please-Ignore": "true" },
},
});
const llm = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.0 });
// complete api
const response1 = await llm.complete("How are you?");
+1 -1
View File
@@ -1,5 +1,5 @@
{
"version": "0.0.20",
"version": "0.0.21",
"private": true,
"name": "simple",
"dependencies": {
@@ -1,9 +1,9 @@
import {
Document,
ListIndex,
ListRetrieverMode,
serviceContextFromDefaults,
SimpleNodeParser,
SummaryIndex,
SummaryRetrieverMode,
serviceContextFromDefaults,
} from "llamaindex";
import essay from "./essay";
@@ -14,9 +14,11 @@ async function main() {
}),
});
const document = new Document({ text: essay, id_: "essay" });
const index = await ListIndex.fromDocuments([document], { serviceContext });
const index = await SummaryIndex.fromDocuments([document], {
serviceContext,
});
const queryEngine = index.asQueryEngine({
retriever: index.asRetriever({ mode: ListRetrieverMode.LLM }),
retriever: index.asRetriever({ mode: SummaryRetrieverMode.LLM }),
});
const response = await queryEngine.query(
"What did the author do growing up?",
+1 -1
View File
@@ -12,7 +12,7 @@ async function main() {
const document = new Document({ text: essay, id_: "essay" });
const serviceContext = serviceContextFromDefaults({
llm: new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.0 }),
llm: new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.1 }),
});
const index = await VectorStoreIndex.fromDocuments([document], {
+1 -4
View File
@@ -4,7 +4,6 @@ import {
PapaCSVReader,
ResponseSynthesizer,
serviceContextFromDefaults,
SimplePrompt,
VectorStoreIndex,
} from "llamaindex";
@@ -23,9 +22,7 @@ async function main() {
serviceContext,
});
const csvPrompt: SimplePrompt = (input) => {
const { context = "", query = "" } = input;
const csvPrompt = ({ context = "", query = "" }) => {
return `The following CSV file is loaded from ${path}
\`\`\`csv
${context}
+4 -2
View File
@@ -2,12 +2,14 @@ import { OpenAI } from "llamaindex";
(async () => {
const llm = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.0 });
// 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" }]);
const response2 = await llm.chat([
{ content: "Tell me a joke!", role: "user" },
]);
console.log(response2.message.content);
})();
@@ -1,9 +1,9 @@
import {
Document,
ListIndex,
ListRetrieverMode,
serviceContextFromDefaults,
SimpleNodeParser,
SummaryIndex,
SummaryRetrieverMode,
serviceContextFromDefaults,
} from "llamaindex";
import essay from "./essay";
@@ -14,9 +14,11 @@ async function main() {
}),
});
const document = new Document({ text: essay, id_: "essay" });
const index = await ListIndex.fromDocuments([document], { serviceContext });
const index = await SummaryIndex.fromDocuments([document], {
serviceContext,
});
const queryEngine = index.asQueryEngine({
retriever: index.asRetriever({ mode: ListRetrieverMode.LLM }),
retriever: index.asRetriever({ mode: SummaryRetrieverMode.LLM }),
});
const response = await queryEngine.query(
"What did the author do growing up?",
+1 -1
View File
@@ -12,7 +12,7 @@ async function main() {
const document = new Document({ text: essay, id_: "essay" });
const serviceContext = serviceContextFromDefaults({
llm: new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.0 }),
llm: new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.1 }),
});
const index = await VectorStoreIndex.fromDocuments([document], {
+7
View File
@@ -1,5 +1,12 @@
# llamaindex
## 0.0.23
### Patch Changes
- Added MetadataMode to ResponseSynthesizer (thanks @TomPenguin)
- 9d6b2ed: Added Markdown Reader (huge shoutout to @swk777)
## 0.0.22
### Patch Changes
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "llamaindex",
"version": "0.0.22",
"version": "0.0.23",
"dependencies": {
"@anthropic-ai/sdk": "^0.6.1",
"lodash": "^4.17.21",
+18 -13
View File
@@ -1,17 +1,18 @@
import { ChatMessage, OpenAI, ChatResponse, LLM } from "./llm/LLM";
import { v4 as uuidv4 } from "uuid";
import { TextNode } from "./Node";
import {
SimplePrompt,
contextSystemPrompt,
CondenseQuestionPrompt,
ContextSystemPrompt,
defaultCondenseQuestionPrompt,
defaultContextSystemPrompt,
messagesToHistoryStr,
} from "./Prompt";
import { BaseQueryEngine } from "./QueryEngine";
import { Response } from "./Response";
import { BaseRetriever } from "./Retriever";
import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext";
import { v4 as uuidv4 } from "uuid";
import { Event } from "./callbacks/CallbackManager";
import { ChatMessage, LLM, OpenAI } from "./llm/LLM";
/**
* A ChatEngine is used to handle back and forth chats between the application and the LLM.
@@ -70,13 +71,13 @@ export class CondenseQuestionChatEngine implements ChatEngine {
queryEngine: BaseQueryEngine;
chatHistory: ChatMessage[];
serviceContext: ServiceContext;
condenseMessagePrompt: SimplePrompt;
condenseMessagePrompt: CondenseQuestionPrompt;
constructor(init: {
queryEngine: BaseQueryEngine;
chatHistory: ChatMessage[];
serviceContext?: ServiceContext;
condenseMessagePrompt?: SimplePrompt;
condenseMessagePrompt?: CondenseQuestionPrompt;
}) {
this.queryEngine = init.queryEngine;
this.chatHistory = init?.chatHistory ?? [];
@@ -92,14 +93,14 @@ export class CondenseQuestionChatEngine implements ChatEngine {
return this.serviceContext.llm.complete(
defaultCondenseQuestionPrompt({
question: question,
chat_history: chatHistoryStr,
})
chatHistory: chatHistoryStr,
}),
);
}
async chat(
message: string,
chatHistory?: ChatMessage[] | undefined
chatHistory?: ChatMessage[] | undefined,
): Promise<Response> {
chatHistory = chatHistory ?? this.chatHistory;
@@ -129,16 +130,20 @@ export class ContextChatEngine implements ChatEngine {
retriever: BaseRetriever;
chatModel: OpenAI;
chatHistory: ChatMessage[];
contextSystemPrompt: ContextSystemPrompt;
constructor(init: {
retriever: BaseRetriever;
chatModel?: OpenAI;
chatHistory?: ChatMessage[];
contextSystemPrompt?: ContextSystemPrompt;
}) {
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;
}
async chat(message: string, chatHistory?: ChatMessage[] | undefined) {
@@ -151,11 +156,11 @@ export class ContextChatEngine implements ChatEngine {
};
const sourceNodesWithScore = await this.retriever.retrieve(
message,
parentEvent
parentEvent,
);
const systemMessage: ChatMessage = {
content: contextSystemPrompt({
content: this.contextSystemPrompt({
context: sourceNodesWithScore
.map((r) => (r.node as TextNode).text)
.join("\n\n"),
@@ -167,7 +172,7 @@ export class ContextChatEngine implements ChatEngine {
const response = await this.chatModel.chat(
[systemMessage, ...chatHistory],
parentEvent
parentEvent,
);
chatHistory.push(response.message);
@@ -175,7 +180,7 @@ export class ContextChatEngine implements ChatEngine {
return new Response(
response.message.content,
sourceNodesWithScore.map((r) => r.node)
sourceNodesWithScore.map((r) => r.node),
);
}
+30 -23
View File
@@ -22,9 +22,7 @@ DEFAULT_TEXT_QA_PROMPT_TMPL = (
)
*/
export const defaultTextQaPrompt: SimplePrompt = (input) => {
const { context = "", query = "" } = input;
export const defaultTextQaPrompt = ({ context = "", query = "" }) => {
return `Context information is below.
---------------------
${context}
@@ -34,6 +32,8 @@ Query: ${query}
Answer:`;
};
export type TextQaPrompt = typeof defaultTextQaPrompt;
/*
DEFAULT_SUMMARY_PROMPT_TMPL = (
"Write a summary of the following. Try to use only the "
@@ -48,9 +48,7 @@ DEFAULT_SUMMARY_PROMPT_TMPL = (
)
*/
export const defaultSummaryPrompt: SimplePrompt = (input) => {
const { context = "" } = input;
export const defaultSummaryPrompt = ({ context = "" }) => {
return `Write a summary of the following. Try to use only the information provided. Try to include as many key details as possible.
@@ -61,6 +59,8 @@ SUMMARY:"""
`;
};
export type SummaryPrompt = typeof defaultSummaryPrompt;
/*
DEFAULT_REFINE_PROMPT_TMPL = (
"The original query is as follows: {query_str}\n"
@@ -77,9 +77,11 @@ DEFAULT_REFINE_PROMPT_TMPL = (
)
*/
export const defaultRefinePrompt: SimplePrompt = (input) => {
const { query = "", existingAnswer = "", context = "" } = input;
export const defaultRefinePrompt = ({
query = "",
existingAnswer = "",
context = "",
}) => {
return `The original query is as follows: ${query}
We have provided an existing answer: ${existingAnswer}
We have the opportunity to refine the existing answer (only if needed) with some more context below.
@@ -90,6 +92,8 @@ Given the new context, refine the original answer to better answer the query. If
Refined Answer:`;
};
export type RefinePrompt = typeof defaultRefinePrompt;
/*
DEFAULT_TREE_SUMMARIZE_TMPL = (
"Context information from multiple sources is below.\n"
@@ -103,9 +107,7 @@ DEFAULT_TREE_SUMMARIZE_TMPL = (
)
*/
export const defaultTreeSummarizePrompt: SimplePrompt = (input) => {
const { context = "", query = "" } = input;
export const defaultTreeSummarizePrompt = ({ context = "", query = "" }) => {
return `Context information from multiple sources is below.
---------------------
${context}
@@ -115,9 +117,9 @@ Query: ${query}
Answer:`;
};
export const defaultChoiceSelectPrompt: SimplePrompt = (input) => {
const { context = "", query = "" } = input;
export type TreeSummarizePrompt = typeof defaultTreeSummarizePrompt;
export const defaultChoiceSelectPrompt = ({ context = "", query = "" }) => {
return `A list of documents is shown below. Each document has a number next to it along
with a summary of the document. A question is also provided.
Respond with the numbers of the documents
@@ -149,6 +151,8 @@ Question: ${query}
Answer:`;
};
export type ChoiceSelectPrompt = typeof defaultChoiceSelectPrompt;
/*
PREFIX = """\
Given a user question, and a list of tools, output a list of relevant sub-questions \
@@ -266,9 +270,7 @@ const exampleOutput: SubQuestion[] = [
},
];
export const defaultSubQuestionPrompt: SimplePrompt = (input) => {
const { toolsStr, queryStr } = input;
export const defaultSubQuestionPrompt = ({ toolsStr = "", queryStr = "" }) => {
return `Given a user question, and a list of tools, output a list of relevant sub-questions that when composed can help answer the full user question:
# Example 1
@@ -298,6 +300,8 @@ ${queryStr}
`;
};
export type SubQuestionPrompt = typeof defaultSubQuestionPrompt;
// DEFAULT_TEMPLATE = """\
// Given a conversation (between Human and Assistant) and a follow up message from Human, \
// rewrite the message to be a standalone question that captures all relevant context \
@@ -312,9 +316,10 @@ ${queryStr}
// <Standalone question>
// """
export const defaultCondenseQuestionPrompt: SimplePrompt = (input) => {
const { chatHistory, question } = input;
export const defaultCondenseQuestionPrompt = ({
chatHistory = "",
question = "",
}) => {
return `Given a conversation (between Human and Assistant) and a follow up message from Human, rewrite the message to be a standalone question that captures all relevant context from the conversation.
<Chat History>
@@ -327,6 +332,8 @@ ${question}
`;
};
export type CondenseQuestionPrompt = typeof defaultCondenseQuestionPrompt;
export function messagesToHistoryStr(messages: ChatMessage[]) {
return messages.reduce((acc, message) => {
acc += acc ? "\n" : "";
@@ -339,11 +346,11 @@ export function messagesToHistoryStr(messages: ChatMessage[]) {
}, "");
}
export const contextSystemPrompt: SimplePrompt = (input) => {
const { context } = input;
export const defaultContextSystemPrompt = ({ context = "" }) => {
return `Context information is below.
---------------------
${context}
---------------------`;
};
export type ContextSystemPrompt = typeof defaultContextSystemPrompt;
+3 -3
View File
@@ -4,7 +4,7 @@ import {
SubQuestionOutputParser,
} from "./OutputParser";
import {
SimplePrompt,
SubQuestionPrompt,
buildToolsText,
defaultSubQuestionPrompt,
} from "./Prompt";
@@ -28,7 +28,7 @@ export interface BaseQuestionGenerator {
*/
export class LLMQuestionGenerator implements BaseQuestionGenerator {
llm: LLM;
prompt: SimplePrompt;
prompt: SubQuestionPrompt;
outputParser: BaseOutputParser<StructuredOutput<SubQuestion[]>>;
constructor(init?: Partial<LLMQuestionGenerator>) {
@@ -45,7 +45,7 @@ export class LLMQuestionGenerator implements BaseQuestionGenerator {
this.prompt({
toolsStr,
queryStr,
})
}),
)
).message.content;
+24 -14
View File
@@ -1,6 +1,9 @@
import { MetadataMode, NodeWithScore } from "./Node";
import {
RefinePrompt,
SimplePrompt,
TextQaPrompt,
TreeSummarizePrompt,
defaultRefinePrompt,
defaultTextQaPrompt,
defaultTreeSummarizePrompt,
@@ -73,13 +76,13 @@ export class SimpleResponseBuilder implements BaseResponseBuilder {
*/
export class Refine implements BaseResponseBuilder {
serviceContext: ServiceContext;
textQATemplate: SimplePrompt;
refineTemplate: SimplePrompt;
textQATemplate: TextQaPrompt;
refineTemplate: RefinePrompt;
constructor(
serviceContext: ServiceContext,
textQATemplate?: SimplePrompt,
refineTemplate?: SimplePrompt,
textQATemplate?: TextQaPrompt,
refineTemplate?: RefinePrompt,
) {
this.serviceContext = serviceContext;
this.textQATemplate = textQATemplate ?? defaultTextQaPrompt;
@@ -209,9 +212,14 @@ export class CompactAndRefine extends Refine {
*/
export class TreeSummarize implements BaseResponseBuilder {
serviceContext: ServiceContext;
summaryTemplate: TreeSummarizePrompt;
constructor(serviceContext: ServiceContext) {
constructor(
serviceContext: ServiceContext,
summaryTemplate?: TreeSummarizePrompt,
) {
this.serviceContext = serviceContext;
this.summaryTemplate = summaryTemplate ?? defaultTreeSummarizePrompt;
}
async getResponse(
@@ -219,21 +227,19 @@ export class TreeSummarize implements BaseResponseBuilder {
textChunks: string[],
parentEvent?: Event,
): Promise<string> {
const summaryTemplate: SimplePrompt = defaultTreeSummarizePrompt;
if (!textChunks || textChunks.length === 0) {
throw new Error("Must have at least one text chunk");
}
const packedTextChunks = this.serviceContext.promptHelper.repack(
summaryTemplate,
this.summaryTemplate,
textChunks,
);
if (packedTextChunks.length === 1) {
return (
await this.serviceContext.llm.complete(
summaryTemplate({
this.summaryTemplate({
context: packedTextChunks[0],
}),
parentEvent,
@@ -243,7 +249,7 @@ export class TreeSummarize implements BaseResponseBuilder {
const summaries = await Promise.all(
packedTextChunks.map((chunk) =>
this.serviceContext.llm.complete(
summaryTemplate({
this.summaryTemplate({
context: chunk,
}),
parentEvent,
@@ -298,9 +304,13 @@ export class ResponseSynthesizer {
this.metadataMode = metadataMode;
}
async synthesize(query: string, nodes: NodeWithScore[], parentEvent?: Event) {
let textChunks: string[] = nodes.map((node) =>
node.node.getContent(this.metadataMode)
async synthesize(
query: string,
nodesWithScore: NodeWithScore[],
parentEvent?: Event,
) {
let textChunks: string[] = nodesWithScore.map(({ node }) =>
node.getContent(this.metadataMode),
);
const response = await this.responseBuilder.getResponse(
query,
@@ -309,7 +319,7 @@ export class ResponseSynthesizer {
);
return new Response(
response,
nodes.map((node) => node.node),
nodesWithScore.map(({ node }) => node),
);
}
}
+1 -1
View File
@@ -1,3 +1,3 @@
export * from "./BaseIndex";
export * from "./list";
export * from "./summary";
export * from "./vectorStore";
-5
View File
@@ -1,5 +0,0 @@
export { ListIndex, ListRetrieverMode } from "./ListIndex";
export {
ListIndexRetriever,
ListIndexLLMRetriever,
} from "./ListIndexRetriever";
@@ -10,11 +10,11 @@ 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,
@@ -22,17 +22,17 @@ import {
IndexStructType,
} from "../BaseIndex";
import {
ListIndexLLMRetriever,
ListIndexRetriever,
} from "./ListIndexRetriever";
SummaryIndexLLMRetriever,
SummaryIndexRetriever,
} from "./SummaryIndexRetriever";
export enum ListRetrieverMode {
export enum SummaryRetrieverMode {
DEFAULT = "default",
// EMBEDDING = "embedding",
LLM = "llm",
}
export interface ListIndexOptions {
export interface SummaryIndexOptions {
nodes?: BaseNode[];
indexStruct?: IndexList;
indexId?: string;
@@ -41,14 +41,14 @@ export interface ListIndexOptions {
}
/**
* A ListIndex keeps nodes in a sequential list structure
* A SummaryIndex keeps nodes in a sequential order for use with summarization.
*/
export class ListIndex extends BaseIndex<IndexList> {
export class SummaryIndex extends BaseIndex<IndexList> {
constructor(init: BaseIndexInit<IndexList>) {
super(init);
}
static async init(options: ListIndexOptions): Promise<ListIndex> {
static async init(options: SummaryIndexOptions): Promise<SummaryIndex> {
const storageContext =
options.storageContext ?? (await storageContextFromDefaults({}));
const serviceContext =
@@ -80,23 +80,23 @@ export class ListIndex extends BaseIndex<IndexList> {
// check indexStruct type
if (indexStruct && indexStruct.type !== IndexStructType.LIST) {
throw new Error(
"Attempting to initialize ListIndex with non-list indexStruct",
"Attempting to initialize SummaryIndex with non-list indexStruct",
);
}
if (indexStruct) {
if (options.nodes) {
throw new Error(
"Cannot initialize VectorStoreIndex with both nodes and indexStruct",
"Cannot initialize SummaryIndex with both nodes and indexStruct",
);
}
} else {
if (!options.nodes) {
throw new Error(
"Cannot initialize VectorStoreIndex without nodes or indexStruct",
"Cannot initialize SummaryIndex without nodes or indexStruct",
);
}
indexStruct = await ListIndex.buildIndexFromNodes(
indexStruct = await SummaryIndex.buildIndexFromNodes(
options.nodes,
storageContext.docStore,
);
@@ -104,7 +104,7 @@ export class ListIndex extends BaseIndex<IndexList> {
await indexStore.addIndexStruct(indexStruct);
}
return new ListIndex({
return new SummaryIndex({
storageContext,
serviceContext,
docStore,
@@ -119,7 +119,7 @@ export class ListIndex extends BaseIndex<IndexList> {
storageContext?: StorageContext;
serviceContext?: ServiceContext;
} = {},
): Promise<ListIndex> {
): Promise<SummaryIndex> {
let { storageContext, serviceContext } = args;
storageContext = storageContext ?? (await storageContextFromDefaults({}));
serviceContext = serviceContext ?? serviceContextFromDefaults({});
@@ -131,7 +131,7 @@ export class ListIndex extends BaseIndex<IndexList> {
}
const nodes = serviceContext.nodeParser.getNodesFromDocuments(documents);
const index = await ListIndex.init({
const index = await SummaryIndex.init({
nodes,
storageContext,
serviceContext,
@@ -139,14 +139,14 @@ export class ListIndex extends BaseIndex<IndexList> {
return index;
}
asRetriever(options?: { mode: ListRetrieverMode }): BaseRetriever {
const { mode = ListRetrieverMode.DEFAULT } = options ?? {};
asRetriever(options?: { mode: SummaryRetrieverMode }): BaseRetriever {
const { mode = SummaryRetrieverMode.DEFAULT } = options ?? {};
switch (mode) {
case ListRetrieverMode.DEFAULT:
return new ListIndexRetriever(this);
case ListRetrieverMode.LLM:
return new ListIndexLLMRetriever(this);
case SummaryRetrieverMode.DEFAULT:
return new SummaryIndexRetriever(this);
case SummaryRetrieverMode.LLM:
return new SummaryIndexLLMRetriever(this);
default:
throw new Error(`Unknown retriever mode: ${mode}`);
}
@@ -253,4 +253,5 @@ export class ListIndex extends BaseIndex<IndexList> {
}
// Legacy
export type GPTListIndex = ListIndex;
export type ListIndex = SummaryIndex;
export type ListRetrieverMode = SummaryRetrieverMode;
@@ -1,25 +1,25 @@
import { BaseRetriever } from "../../Retriever";
import _ from "lodash";
import { globalsHelper } from "../../GlobalsHelper";
import { NodeWithScore } from "../../Node";
import { ListIndex } from "./ListIndex";
import { ChoiceSelectPrompt, defaultChoiceSelectPrompt } from "../../Prompt";
import { BaseRetriever } from "../../Retriever";
import { ServiceContext } from "../../ServiceContext";
import { Event } from "../../callbacks/CallbackManager";
import { SummaryIndex } from "./SummaryIndex";
import {
NodeFormatterFunction,
ChoiceSelectParserFunction,
NodeFormatterFunction,
defaultFormatNodeBatchFn,
defaultParseChoiceSelectAnswerFn,
} from "./utils";
import { SimplePrompt, defaultChoiceSelectPrompt } from "../../Prompt";
import _ from "lodash";
import { globalsHelper } from "../../GlobalsHelper";
import { Event } from "../../callbacks/CallbackManager";
/**
* Simple retriever for ListIndex that returns all nodes
* Simple retriever for SummaryIndex that returns all nodes
*/
export class ListIndexRetriever implements BaseRetriever {
index: ListIndex;
export class SummaryIndexRetriever implements BaseRetriever {
index: SummaryIndex;
constructor(index: ListIndex) {
constructor(index: SummaryIndex) {
this.index = index;
}
@@ -51,23 +51,23 @@ export class ListIndexRetriever implements BaseRetriever {
}
/**
* LLM retriever for ListIndex.
* LLM retriever for SummaryIndex which lets you select the most relevant chunks.
*/
export class ListIndexLLMRetriever implements BaseRetriever {
index: ListIndex;
choiceSelectPrompt: SimplePrompt;
export class SummaryIndexLLMRetriever implements BaseRetriever {
index: SummaryIndex;
choiceSelectPrompt: ChoiceSelectPrompt;
choiceBatchSize: number;
formatNodeBatchFn: NodeFormatterFunction;
parseChoiceSelectAnswerFn: ChoiceSelectParserFunction;
serviceContext: ServiceContext;
constructor(
index: ListIndex,
choiceSelectPrompt?: SimplePrompt,
index: SummaryIndex,
choiceSelectPrompt?: ChoiceSelectPrompt,
choiceBatchSize: number = 10,
formatNodeBatchFn?: NodeFormatterFunction,
parseChoiceSelectAnswerFn?: ChoiceSelectParserFunction,
serviceContext?: ServiceContext
serviceContext?: ServiceContext,
) {
this.index = index;
this.choiceSelectPrompt = choiceSelectPrompt || defaultChoiceSelectPrompt;
@@ -95,7 +95,7 @@ export class ListIndexLLMRetriever implements BaseRetriever {
// parseResult is a map from doc number to relevance score
const parseResult = this.parseChoiceSelectAnswerFn(
rawResponse,
nodesBatch.length
nodesBatch.length,
);
const choiceNodeIds = nodeIdsBatch.filter((nodeId, idx) => {
return `${idx}` in parseResult;
@@ -128,3 +128,7 @@ export class ListIndexLLMRetriever implements BaseRetriever {
return this.serviceContext;
}
}
// Legacy
export type ListIndexRetriever = SummaryIndexRetriever;
export type ListIndexLLMRetriever = SummaryIndexLLMRetriever;
@@ -0,0 +1,10 @@
export { SummaryIndex, SummaryRetrieverMode } from "./SummaryIndex";
export type { ListIndex, ListRetrieverMode } from "./SummaryIndex";
export {
SummaryIndexLLMRetriever,
SummaryIndexRetriever,
} from "./SummaryIndexRetriever";
export type {
ListIndexLLMRetriever,
ListIndexRetriever,
} from "./SummaryIndexRetriever";
+17 -19
View File
@@ -1,18 +1,18 @@
import { VectorStoreIndex } from "../indices/vectorStore/VectorStoreIndex";
import { OpenAIEmbedding } from "../Embedding";
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 { ListIndex, ListRetrieverMode } from "../indices/list";
import {
ResponseSynthesizer,
SimpleResponseBuilder,
} from "../ResponseSynthesizer";
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
@@ -65,10 +65,9 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
});
test("For VectorStoreIndex w/ a SimpleResponseBuilder", async () => {
const vectorStoreIndex = await VectorStoreIndex.fromDocuments(
[document],
{ serviceContext }
);
const vectorStoreIndex = await VectorStoreIndex.fromDocuments([document], {
serviceContext,
});
const queryEngine = vectorStoreIndex.asQueryEngine();
const query = "What is the author's name?";
const response = await queryEngine.query(query);
@@ -132,21 +131,20 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
// both retrieval and streaming should have
// the same parent event
expect(streamCallbackData[0].event.parentId).toBe(
retrieveCallbackData[0].event.parentId
retrieveCallbackData[0].event.parentId,
);
});
test("For ListIndex w/ a ListIndexRetriever", async () => {
const listIndex = await ListIndex.fromDocuments(
[document],
{ serviceContext },
);
test("For SummaryIndex w/ a SummaryIndexRetriever", async () => {
const summaryIndex = await SummaryIndex.fromDocuments([document], {
serviceContext,
});
const responseBuilder = new SimpleResponseBuilder(serviceContext);
const responseSynthesizer = new ResponseSynthesizer({
serviceContext: serviceContext,
responseBuilder,
});
const queryEngine = listIndex.asQueryEngine({
const queryEngine = summaryIndex.asQueryEngine({
responseSynthesizer,
});
const query = "What is the author's name?";
@@ -211,7 +209,7 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
// both retrieval and streaming should have
// the same parent event
expect(streamCallbackData[0].event.parentId).toBe(
retrieveCallbackData[0].event.parentId
retrieveCallbackData[0].event.parentId,
);
});
});