chore: use Logger for core (#2139)

.changeset/vast-roses-cheer.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+Use logger interface instead of directly hardcoding console.log

packages/core/src/embeddings/base.ts

@@ -1,3 +1,4 @@
+import { consoleLogger, emptyLogger, type Logger } from "@llamaindex/env";
 import type { Tokenizers } from "@llamaindex/env/tokenizers";
 import type { MessageContentDetail } from "../llms";
 import { BaseNode, MetadataMode, TransformComponent } from "../schema";
@@ -18,6 +19,7 @@ export type EmbeddingInfo = {
 export type BaseEmbeddingOptions = {
   logProgress?: boolean;
   progressCallback?: (current: number, total: number) => void;
+  logger?: Logger;
 };
 
 export abstract class BaseEmbedding extends TransformComponent<
@@ -133,6 +135,9 @@ export async function batchEmbeddings<T>(
 
   const curBatch: T[] = [];
 
+  const logger =
+    options?.logger ?? (options?.logProgress ? consoleLogger : emptyLogger);
+
   for (let i = 0; i < queue.length; i++) {
     curBatch.push(queue[i]!);
     if (i == queue.length - 1 || curBatch.length == chunkSize) {
@@ -143,7 +148,7 @@ export async function batchEmbeddings<T>(
         options?.progressCallback?.(i + 1, queue.length);
       }
       if (options?.logProgress) {
-        console.log(`getting embedding progress: ${i + 1} / ${queue.length}`);
+        logger.log(`getting embedding progress: ${i + 1} / ${queue.length}`);
       }
 
       curBatch.length = 0;

packages/core/src/memory/memory.ts

@@ -1,3 +1,4 @@
+import { consoleLogger, type Logger } from "@llamaindex/env";
 import { Settings } from "../global";
 import type { ChatMessage, LLM } from "../llms";
 import { extractText } from "../utils";
@@ -38,6 +39,11 @@ export type MemoryOptions<TMessageOptions extends object = object> = {
    * This default LLM can be overridden by the LLM passed in the `getLLM` method.
    */
   llm?: LLM | undefined;
+
+  /**
+   * Logger for memory operations
+   */
+  logger?: Logger;
 };
 
 export class Memory<
@@ -76,6 +82,10 @@ export class Memory<
    * The default LLM to use for memory retrieval.
    */
   private llm: LLM | undefined;
+  /**
+   * Logger for memory operations
+   */
+  private logger: Logger;
 
   constructor(
     messages: MemoryMessage<TMessageOptions>[] = [],
@@ -87,6 +97,7 @@ export class Memory<
       options.shortTermTokenLimitRatio ?? DEFAULT_SHORT_TERM_TOKEN_LIMIT_RATIO;
     this.memoryBlocks = options.memoryBlocks ?? [];
     this.memoryCursor = options.memoryCursor ?? 0;
+    this.logger = options.logger ?? consoleLogger;
     this.initLLM(options.llm);
 
     this.adapters = {
@@ -309,7 +320,7 @@ export class Memory<
           addedTokenCount += messageTokenCount;
         }
       } catch (error) {
-        console.warn(
+        this.logger.warn(
           `Failed to get content from memory block ${block.id}:`,
           error,
         );
@@ -371,7 +382,7 @@ export class Memory<
       try {
         await block.put(newMessages);
       } catch (error) {
-        console.warn(
+        this.logger.warn(
           `Failed to process messages into memory block ${block.id}:`,
           error,
         );

packages/core/src/node-parser/sentence-splitter.ts

@@ -1,3 +1,4 @@
+import { consoleLogger, type Logger } from "@llamaindex/env";
 import type { Tokenizer } from "@llamaindex/env/tokenizers";
 import { z } from "zod";
 import { Settings } from "../global";
@@ -48,9 +49,11 @@ export class SentenceSplitter extends MetadataAwareTextSplitter {
   #splitFns: Set<TextSplitterFn> = new Set();
   #subSentenceSplitFns: Set<TextSplitterFn> = new Set();
   #tokenizer: Tokenizer;
+  #logger: Logger;
 
   constructor(
-    params?: z.input<typeof sentenceSplitterSchema> & SplitterParams,
+    params?: z.input<typeof sentenceSplitterSchema> &
+      SplitterParams & { logger?: Logger },
   ) {
     super();
     if (params) {
@@ -66,6 +69,7 @@ export class SentenceSplitter extends MetadataAwareTextSplitter {
       this.extraAbbreviations,
     );
     this.#tokenizer = params?.tokenizer ?? Settings.tokenizer;
+    this.#logger = params?.logger ?? consoleLogger;
     this.#splitFns.add(splitBySep(this.paragraphSeparator));
     this.#splitFns.add(this.#chunkingTokenizerFn);
 
@@ -82,7 +86,7 @@ export class SentenceSplitter extends MetadataAwareTextSplitter {
         `Metadata length (${metadataLength}) is longer than chunk size (${this.chunkSize}). Consider increasing the chunk size or decreasing the size of your metadata to avoid this.`,
       );
     } else if (effectiveChunkSize < 50) {
-      console.log(
+      this.#logger.log(
         `Metadata length (${metadataLength}) is close to chunk size (${this.chunkSize}). Resulting chunks are less than 50 tokens. Consider increasing the chunk size or decreasing the size of your metadata to avoid this.`,
       );
     }

packages/core/src/node-parser/token-text-splitter.ts

@@ -1,3 +1,4 @@
+import { consoleLogger, type Logger } from "@llamaindex/env";
 import type { Tokenizer } from "@llamaindex/env/tokenizers";
 import { z } from "zod";
 import { DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE, Settings } from "../global";
@@ -21,9 +22,11 @@ export class TokenTextSplitter extends MetadataAwareTextSplitter {
   backupSeparators: string[] = ["\n"];
   #tokenizer: Tokenizer;
   #splitFns: Array<(text: string) => string[]> = [];
+  #logger: Logger;
 
   constructor(
-    params?: SplitterParams & Partial<z.infer<typeof tokenTextSplitterSchema>>,
+    params?: SplitterParams &
+      Partial<z.infer<typeof tokenTextSplitterSchema>> & { logger?: Logger },
   ) {
     super();
 
@@ -42,6 +45,7 @@ export class TokenTextSplitter extends MetadataAwareTextSplitter {
     }
 
     this.#tokenizer = params?.tokenizer ?? Settings.tokenizer;
+    this.#logger = params?.logger ?? consoleLogger;
 
     const allSeparators = [this.separator, ...this.backupSeparators];
     this.#splitFns = allSeparators.map((sep) => splitBySep(sep));
@@ -65,7 +69,7 @@ export class TokenTextSplitter extends MetadataAwareTextSplitter {
           `Consider increasing the chunk size or decreasing the size of your metadata to avoid this.`,
       );
     } else if (effectiveChunkSize < 50) {
-      console.warn(
+      this.#logger.warn(
         `Metadata length (${metadataLength}) is close to chunk size (${this.chunkSize}). ` +
           `Resulting chunks are less than 50 tokens. Consider increasing the chunk size or decreasing the size of your metadata to avoid this.`,
       );
@@ -148,7 +152,7 @@ export class TokenTextSplitter extends MetadataAwareTextSplitter {
       const splitLength = this.tokenSize(split);
 
       if (splitLength > chunkSize) {
-        console.warn(
+        this.#logger.warn(
           `Got a split of size ${splitLength}, larger than chunk size ${chunkSize}.`,
         );
       }

packages/core/src/storage/doc-store/kv-document-store.ts

@@ -1,3 +1,4 @@
+import { consoleLogger, type Logger } from "@llamaindex/env";
 import { DEFAULT_NAMESPACE } from "../../global";
 import { BaseNode, ObjectType, type StoredValue } from "../../schema";
 import type { BaseKVStore } from "../kv-store";
@@ -16,13 +17,19 @@ export class KVDocumentStore extends BaseDocumentStore {
   private nodeCollection: string;
   private refDocCollection: string;
   private metadataCollection: string;
+  private logger: Logger;
 
-  constructor(kvstore: BaseKVStore, namespace: string = DEFAULT_NAMESPACE) {
+  constructor(
+    kvstore: BaseKVStore,
+    namespace: string = DEFAULT_NAMESPACE,
+    options?: { logger?: Logger },
+  ) {
     super();
     this.kvstore = kvstore;
     this.nodeCollection = `${namespace}/data`;
     this.refDocCollection = `${namespace}/ref_doc_info`;
     this.metadataCollection = `${namespace}/metadata`;
+    this.logger = options?.logger ?? consoleLogger;
   }
 
   async docs(): Promise<Record<string, BaseNode>> {
@@ -33,7 +40,7 @@ export class KVDocumentStore extends BaseDocumentStore {
       if (isValidDocJson(value)) {
         docs[key] = jsonToDoc(value, this.serializer);
       } else {
-        console.warn(`Invalid JSON for docId ${key}`);
+        this.logger.warn(`Invalid JSON for docId ${key}`);
       }
     }
     return docs;

packages/core/src/storage/index-store/index.ts

@@ -1,4 +1,4 @@
-import { path } from "@llamaindex/env";
+import { path, type Logger } from "@llamaindex/env";
 import { IndexStruct, jsonToIndexStruct } from "../../data-structs";
 import {
   DEFAULT_INDEX_STORE_PERSIST_FILENAME,
@@ -8,8 +8,8 @@ import {
 import {
   BaseInMemoryKVStore,
   BaseKVStore,
-  type DataType,
   SimpleKVStore,
+  type DataType,
 } from "../kv-store";
 
 export const DEFAULT_PERSIST_PATH = path.join(
@@ -84,16 +84,23 @@ export class SimpleIndexStore extends KVIndexStore {
 
   static async fromPersistDir(
     persistDir: string = DEFAULT_PERSIST_DIR,
+    options?: { logger?: Logger },
   ): Promise<SimpleIndexStore> {
     const persistPath = path.join(
       persistDir,
       DEFAULT_INDEX_STORE_PERSIST_FILENAME,
     );
-    return this.fromPersistPath(persistPath);
+    return this.fromPersistPath(persistPath, options);
   }
 
-  static async fromPersistPath(persistPath: string): Promise<SimpleIndexStore> {
-    const simpleKVStore = await SimpleKVStore.fromPersistPath(persistPath);
+  static async fromPersistPath(
+    persistPath: string,
+    options?: { logger?: Logger },
+  ): Promise<SimpleIndexStore> {
+    const simpleKVStore = await SimpleKVStore.fromPersistPath(
+      persistPath,
+      options,
+    );
     return new SimpleIndexStore(simpleKVStore);
   }
 

packages/core/src/storage/kv-store/index.ts

@@ -1,4 +1,4 @@
-import { fs, path } from "@llamaindex/env";
+import { consoleLogger, fs, path, type Logger } from "@llamaindex/env";
 
 import { DEFAULT_COLLECTION } from "../../global";
 import type { StoredValue } from "../../schema";
@@ -98,7 +98,11 @@ export class SimpleKVStore extends BaseKVStore {
     await fs.writeFile(persistPath, JSON.stringify(this.data));
   }
 
-  static async fromPersistPath(persistPath: string): Promise<SimpleKVStore> {
+  static async fromPersistPath(
+    persistPath: string,
+    options?: { logger?: Logger },
+  ): Promise<SimpleKVStore> {
+    const logger = options?.logger ?? consoleLogger;
     const dirPath = path.dirname(persistPath);
     if (!(await exists(dirPath))) {
       await fs.mkdir(dirPath, { recursive: true });
@@ -106,7 +110,7 @@ export class SimpleKVStore extends BaseKVStore {
 
     let data: DataType = {};
     if (!(await exists(persistPath))) {
-      console.info(`Starting new store from path: ${persistPath}`);
+      logger.log(`Starting new store from path: ${persistPath}`);
     } else {
       try {
         const fileData = await fs.readFile(persistPath);

packages/core/src/tools/function-tool.ts

@@ -1,3 +1,4 @@
+import { consoleLogger, type Logger } from "@llamaindex/env";
 import type { JSONSchemaType } from "ajv";
 import { z } from "zod";
 import { zodToJsonSchema } from "zod-to-json-schema";
@@ -14,11 +15,13 @@ export class FunctionTool<
   #additionalArg: AdditionalToolArgument | undefined;
   readonly #metadata: ToolMetadata<JSONSchemaType<T>>;
   readonly #zodType: z.ZodType<T> | null = null;
+  readonly #logger: Logger;
   constructor(
     fn: (input: T, additionalArg?: AdditionalToolArgument) => R,
     metadata: ToolMetadata<JSONSchemaType<T>>,
     zodType?: z.ZodType<T>,
     additionalArg?: AdditionalToolArgument,
+    logger?: Logger,
   ) {
     this.#fn = fn;
     this.#metadata = metadata;
@@ -26,6 +29,7 @@ export class FunctionTool<
       this.#zodType = zodType;
     }
     this.#additionalArg = additionalArg;
+    this.#logger = logger ?? consoleLogger;
   }
 
   static from<T, AdditionalToolArgument extends object = object>(
@@ -140,7 +144,7 @@ export class FunctionTool<
       if (result.success) {
         params = result.data;
       } else {
-        console.warn(result.error.errors);
+        this.#logger.warn(result.error.errors);
       }
     }
     return this.#fn.call(null, params, this.#additionalArg);

packages/llamaindex/src/storage/docStore/SimpleDocumentStore.ts

@@ -8,7 +8,7 @@ import {
   BaseInMemoryKVStore,
   SimpleKVStore,
 } from "@llamaindex/core/storage/kv-store";
-import { path } from "@llamaindex/env";
+import { path, type Logger } from "@llamaindex/env";
 import _ from "lodash";
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -27,19 +27,28 @@ export class SimpleDocumentStore extends KVDocumentStore {
   static async fromPersistDir(
     persistDir: string = DEFAULT_PERSIST_DIR,
     namespace?: string,
+    options?: { logger?: Logger },
   ): Promise<SimpleDocumentStore> {
     const persistPath = path.join(
       persistDir,
       DEFAULT_DOC_STORE_PERSIST_FILENAME,
     );
-    return await SimpleDocumentStore.fromPersistPath(persistPath, namespace);
+    return await SimpleDocumentStore.fromPersistPath(
+      persistPath,
+      namespace,
+      options,
+    );
   }
 
   static async fromPersistPath(
     persistPath: string,
     namespace?: string,
+    options?: { logger?: Logger },
   ): Promise<SimpleDocumentStore> {
-    const simpleKVStore = await SimpleKVStore.fromPersistPath(persistPath);
+    const simpleKVStore = await SimpleKVStore.fromPersistPath(
+      persistPath,
+      options,
+    );
     return new SimpleDocumentStore(simpleKVStore, namespace);
   }
 

packages/llamaindex/src/vector-store/SimpleVectorStore.ts

@@ -18,7 +18,7 @@ import {
   type VectorStoreQuery,
   type VectorStoreQueryResult,
 } from "@llamaindex/core/vector-store";
-import { fs, path } from "@llamaindex/env";
+import { consoleLogger, fs, path, type Logger } from "@llamaindex/env";
 import { exists } from "../storage/FileSystem.js";
 
 const LEARNER_MODES = new Set<VectorStoreQueryMode>([
@@ -139,9 +139,14 @@ export class SimpleVectorStore extends BaseVectorStore {
   static async fromPersistDir(
     persistDir: string = DEFAULT_PERSIST_DIR,
     embedModel?: BaseEmbedding,
+    options?: { logger?: Logger },
   ): Promise<SimpleVectorStore> {
     const persistPath = path.join(persistDir, "vector_store.json");
-    return await SimpleVectorStore.fromPersistPath(persistPath, embedModel);
+    return await SimpleVectorStore.fromPersistPath(
+      persistPath,
+      embedModel,
+      options,
+    );
   }
 
   client() {
@@ -273,7 +278,9 @@ export class SimpleVectorStore extends BaseVectorStore {
   static async fromPersistPath(
     persistPath: string,
     embedModel?: BaseEmbedding,
+    options?: { logger?: Logger },
   ): Promise<SimpleVectorStore> {
+    const logger = options?.logger ?? consoleLogger;
     const dirPath = path.dirname(persistPath);
     if (!(await exists(dirPath))) {
       await fs.mkdir(dirPath, { recursive: true });
@@ -281,7 +288,7 @@ export class SimpleVectorStore extends BaseVectorStore {
 
     let dataDict: Record<string, unknown> = {};
     if (!(await exists(persistPath))) {
-      console.info(`Starting new store from path: ${persistPath}`);
+      logger.log(`Starting new store from path: ${persistPath}`);
     } else {
       try {
         const fileData = await fs.readFile(persistPath);

packages/llamaindex/tests/storage/StorageContext.test.ts

@@ -47,22 +47,31 @@ describe("StorageContext", () => {
 
   test("persists and loads", async () => {
     const doc = new Document({ text: "test document" });
-    const consoleInfoSpy = vi
-      .spyOn(console, "info")
-      .mockImplementation(() => {});
+    // Create a Logger that spies on log (info) calls
+    const spyLogger = {
+      log: vi.fn(),
+      error: vi.fn(),
+      warn: vi.fn(),
+    };
 
     // storage context from individual stores
     const storageContext = await storageContextFromDefaults({
-      docStore: await SimpleDocumentStore.fromPersistDir(testDir),
-      vectorStore: await SimpleVectorStore.fromPersistDir(testDir),
-      indexStore: await SimpleIndexStore.fromPersistDir(testDir),
+      docStore: await SimpleDocumentStore.fromPersistDir(testDir, undefined, {
+        logger: spyLogger,
+      }),
+      vectorStore: await SimpleVectorStore.fromPersistDir(testDir, undefined, {
+        logger: spyLogger,
+      }),
+      indexStore: await SimpleIndexStore.fromPersistDir(testDir, {
+        logger: spyLogger,
+      }),
     });
 
     const index = await VectorStoreIndex.fromDocuments([doc], {
       storageContext,
     });
-    expect(consoleInfoSpy).toHaveBeenCalledTimes(3);
-    expect(consoleInfoSpy).toHaveBeenCalledWith(
+    expect(spyLogger.log).toHaveBeenCalledTimes(3);
+    expect(spyLogger.log).toHaveBeenCalledWith(
       expect.stringContaining("Starting new store"),
     );
     expect(index).toBeDefined();
@@ -75,13 +84,19 @@ describe("StorageContext", () => {
     // Check that the test data files exist
     await expectTestDataFilesExist(testDir);
 
-    consoleInfoSpy.mockClear();
+    spyLogger.log.mockClear();
 
     // Now, load it again. Since data was persisted, we should not see the error.
     const newStorageContext = await storageContextFromDefaults({
-      docStore: await SimpleDocumentStore.fromPersistDir(testDir),
-      vectorStore: await SimpleVectorStore.fromPersistDir(testDir),
-      indexStore: await SimpleIndexStore.fromPersistDir(testDir),
+      docStore: await SimpleDocumentStore.fromPersistDir(testDir, undefined, {
+        logger: spyLogger,
+      }),
+      vectorStore: await SimpleVectorStore.fromPersistDir(testDir, undefined, {
+        logger: spyLogger,
+      }),
+      indexStore: await SimpleIndexStore.fromPersistDir(testDir, {
+        logger: spyLogger,
+      }),
     });
 
     const loadedIndex = await VectorStoreIndex.init({
@@ -94,9 +109,7 @@ describe("StorageContext", () => {
 
     await expectTestDataFilesExist(testDir);
 
-    expect(consoleInfoSpy).not.toHaveBeenCalled();
-
-    consoleInfoSpy.mockRestore();
+    expect(spyLogger.log).not.toHaveBeenCalled();
   });
 
   test("throws error on corrupted data", async () => {