refactor: move mockLLM to core (#1493)

Co-authored-by: Alex Yang <himself65@outlook.com>

.changeset/perfect-turtles-mate.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+refactor: move mockLLM to core

apps/next/src/app/api/chat/route.ts

@@ -1,9 +1,9 @@
-import { llm } from "@/lib/utils";
+import { MockLLM } from "@llamaindex/core/utils";
 import { LlamaIndexAdapter, type Message } from "ai";
 import { Settings, SimpleChatEngine, type ChatMessage } from "llamaindex";
 import { NextResponse, type NextRequest } from "next/server";
 
-Settings.llm = llm;
+Settings.llm = new MockLLM(); // config your LLM here
 
 export async function POST(request: NextRequest) {
   try {

apps/next/src/components/demo/chat/rsc/ai-action.tsx

@@ -1,5 +1,5 @@
-import { llm } from "@/lib/utils";
 import { Markdown } from "@llamaindex/chat-ui/widgets";
+import { MockLLM } from "@llamaindex/core/utils";
 import { generateId, Message } from "ai";
 import { createAI, createStreamableUI, getMutableAIState } from "ai/rsc";
 import { type ChatMessage, Settings, SimpleChatEngine } from "llamaindex";
@@ -11,7 +11,7 @@ type Actions = {
   chat: (message: Message) => Promise<Message & { display: ReactNode }>;
 };
 
-Settings.llm = llm;
+Settings.llm = new MockLLM(); // config your LLM here
 
 export const AI = createAI<ServerState, FrontendState, Actions>({
   initialAIState: [],

apps/next/src/lib/utils.ts

@@ -1,34 +1,6 @@
-import { clsx, type ClassValue } from "clsx";
-import { LLM, LLMMetadata } from "llamaindex";
-import { twMerge } from "tailwind-merge";
+import { clsx, type ClassValue } from "clsx"
+import { twMerge } from "tailwind-merge"
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs))
 }
-
-class MockLLM  {
-  metadata: LLMMetadata = {
-    model: "MockLLM",
-    temperature: 0.5,
-    topP: 0.5,
-    contextWindow: 1024,
-    tokenizer: undefined,
-  };
-
-  chat() {
-    const mockResponse = "Hello! This is a mock response";
-    return Promise.resolve(
-      new ReadableStream({
-        async start(controller) {
-          for (const char of mockResponse) {
-            controller.enqueue({ delta: char });
-            await new Promise((resolve) => setTimeout(resolve, 20));
-          }
-          controller.close();
-        },
-      }),
-    );
-  }
-}
-
-export const llm = new MockLLM() as unknown as LLM;

packages/core/src/utils/index.ts

@@ -76,6 +76,7 @@ export {
   extractText,
   imageToDataUrl,
   messagesToHistory,
+  MockLLM,
   toToolDescriptions,
 } from "./llms";
 

packages/core/src/utils/llms.ts

@@ -2,6 +2,15 @@ import { fs } from "@llamaindex/env";
 import { filetypemime } from "magic-bytes.js";
 import type {
   ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
+  CompletionResponse,
+  LLM,
+  LLMChatParamsNonStreaming,
+  LLMChatParamsStreaming,
+  LLMCompletionParamsNonStreaming,
+  LLMCompletionParamsStreaming,
+  LLMMetadata,
   MessageContent,
   MessageContentDetail,
   MessageContentTextDetail,
@@ -143,3 +152,82 @@ export async function imageToDataUrl(
   }
   return await blobToDataUrl(input);
 }
+
+export class MockLLM implements LLM {
+  metadata: LLMMetadata;
+  options: {
+    timeBetweenToken: number;
+    responseMessage: string;
+  };
+
+  constructor(options?: {
+    timeBetweenToken?: number;
+    responseMessage?: string;
+    metadata?: LLMMetadata;
+  }) {
+    this.options = {
+      timeBetweenToken: options?.timeBetweenToken ?? 20,
+      responseMessage: options?.responseMessage ?? "This is a mock response",
+    };
+    this.metadata = options?.metadata ?? {
+      model: "MockLLM",
+      temperature: 0.5,
+      topP: 0.5,
+      contextWindow: 1024,
+      tokenizer: undefined,
+    };
+  }
+
+  chat(
+    params: LLMChatParamsStreaming<object, object>,
+  ): Promise<AsyncIterable<ChatResponseChunk>>;
+  chat(
+    params: LLMChatParamsNonStreaming<object, object>,
+  ): Promise<ChatResponse<object>>;
+  async chat(
+    params:
+      | LLMChatParamsStreaming<object, object>
+      | LLMChatParamsNonStreaming<object, object>,
+  ): Promise<AsyncIterable<ChatResponseChunk> | ChatResponse<object>> {
+    const responseMessage = this.options.responseMessage;
+    const timeBetweenToken = this.options.timeBetweenToken;
+
+    if (params.stream) {
+      return (async function* () {
+        for (const char of responseMessage) {
+          yield { delta: char, raw: {} };
+          await new Promise((resolve) => setTimeout(resolve, timeBetweenToken));
+        }
+      })();
+    }
+
+    return {
+      message: { content: responseMessage, role: "assistant" },
+      raw: {},
+    };
+  }
+
+  async complete(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<CompletionResponse>>;
+  async complete(
+    params: LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse>;
+  async complete(
+    params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
+  ): Promise<AsyncIterable<CompletionResponse> | CompletionResponse> {
+    const responseMessage = this.options.responseMessage;
+    const timeBetweenToken = this.options.timeBetweenToken;
+
+    if (params.stream) {
+      return (async function* () {
+        for (const char of responseMessage) {
+          yield { delta: char, text: char, raw: {} };
+          await new Promise((resolve) => setTimeout(resolve, timeBetweenToken));
+        }
+      })();
+    }
+
+    return { text: responseMessage, raw: {} };
+  }
+}