fix: use metadata mode LLM for generating context (run-llama#1133)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>
himself65 · Aug 23, 2024 · eed0b04 · eed0b04
1 parent d9d6c56
commit eed0b04
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/.changeset/tough-cups-doubt.md b/.changeset/tough-cups-doubt.md
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+fix: use LLM metadata mode for generating context of ContextChatEngine
diff --git a/packages/llamaindex/src/engines/chat/ContextChatEngine.ts b/packages/llamaindex/src/engines/chat/ContextChatEngine.ts
@@ -4,7 +4,7 @@ import type {
   MessageContent,
   MessageType,
 } from "@llamaindex/core/llms";
-import { EngineResponse } from "@llamaindex/core/schema";
+import { EngineResponse, MetadataMode } from "@llamaindex/core/schema";
 import {
   extractText,
   streamConverter,
@@ -53,6 +53,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
       contextSystemPrompt: init?.contextSystemPrompt,
       nodePostprocessors: init?.nodePostprocessors,
       contextRole: init?.contextRole,
+      metadataMode: MetadataMode.LLM,
     });
     this.systemPrompt = init.systemPrompt;
   }

diff --git a/packages/llamaindex/src/engines/chat/DefaultContextGenerator.ts b/packages/llamaindex/src/engines/chat/DefaultContextGenerator.ts
@@ -1,5 +1,5 @@
 import type { MessageContent, MessageType } from "@llamaindex/core/llms";
-import { type NodeWithScore } from "@llamaindex/core/schema";
+import { MetadataMode, type NodeWithScore } from "@llamaindex/core/schema";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
 import type { ContextSystemPrompt } from "../../Prompt.js";
 import { defaultContextSystemPrompt } from "../../Prompt.js";
@@ -16,12 +16,14 @@ export class DefaultContextGenerator
   contextSystemPrompt: ContextSystemPrompt;
   nodePostprocessors: BaseNodePostprocessor[];
   contextRole: MessageType;
+  metadataMode?: MetadataMode;
 
   constructor(init: {
     retriever: BaseRetriever;
     contextSystemPrompt?: ContextSystemPrompt;
     nodePostprocessors?: BaseNodePostprocessor[];
     contextRole?: MessageType;
+    metadataMode?: MetadataMode;
   }) {
     super();
 
@@ -30,6 +32,7 @@ export class DefaultContextGenerator
       init?.contextSystemPrompt ?? defaultContextSystemPrompt;
     this.nodePostprocessors = init.nodePostprocessors || [];
     this.contextRole = init.contextRole ?? "system";
+    this.metadataMode = init.metadataMode ?? MetadataMode.NONE;
   }
 
   protected _getPrompts(): { contextSystemPrompt: ContextSystemPrompt } {
@@ -75,6 +78,8 @@ export class DefaultContextGenerator
     const content = await createMessageContent(
       this.contextSystemPrompt,
       nodes.map((r) => r.node),
+      undefined,
+      this.metadataMode,
     );
 
     return {