jan/extensions/assistant-extension/src/index.ts

import {
  fs,
  Assistant,
  MessageRequest,
  events,
  InferenceEngine,
  MessageEvent,
  InferenceEvent,
  joinPath,
  executeOnMain,
  AssistantExtension,
  AssistantEvent,
} from "@janhq/core";

export default class JanAssistantExtension extends AssistantExtension {
  private static readonly _homeDir = "file://assistants";

  controller = new AbortController();
  isCancelled = false;
  retrievalThreadId: string | undefined = undefined;

  async onLoad() {
    // making the assistant directory
    const assistantDirExist = await fs.existsSync(
      JanAssistantExtension._homeDir
    );
    if (
      localStorage.getItem(`${EXTENSION_NAME}-version`) !== VERSION ||
      !assistantDirExist
    ) {
      if (!assistantDirExist)
        await fs.mkdirSync(JanAssistantExtension._homeDir);

      // Write assistant metadata
      await this.createJanAssistant();
      // Finished migration
      localStorage.setItem(`${EXTENSION_NAME}-version`, VERSION);
      // Update the assistant list
      events.emit(AssistantEvent.OnAssistantsUpdate, {});
    }

    // Events subscription
    events.on(MessageEvent.OnMessageSent, (data: MessageRequest) =>
      JanAssistantExtension.handleMessageRequest(data, this)
    );

    events.on(InferenceEvent.OnInferenceStopped, () => {
      JanAssistantExtension.handleInferenceStopped(this);
    });
  }

  private static async handleInferenceStopped(instance: JanAssistantExtension) {
    instance.isCancelled = true;
    instance.controller?.abort();
  }

  private static async handleMessageRequest(
    data: MessageRequest,
    instance: JanAssistantExtension
  ) {
    instance.isCancelled = false;
    instance.controller = new AbortController();

    if (
      data.model?.engine !== InferenceEngine.tool_retrieval_enabled ||
      !data.messages ||
      !data.thread?.assistants[0]?.tools
    ) {
      return;
    }

    const latestMessage = data.messages[data.messages.length - 1];

    // Ingest the document if needed
    if (
      latestMessage &&
      latestMessage.content &&
      typeof latestMessage.content !== "string"
    ) {
      const docFile = latestMessage.content[1]?.doc_url?.url;
      if (docFile) {
        await executeOnMain(
          NODE,
          "toolRetrievalIngestNewDocument",
          docFile,
          data.model?.proxyEngine
        );
      }
    }

    // Load agent on thread changed
    if (instance.retrievalThreadId !== data.threadId) {
      await executeOnMain(NODE, "toolRetrievalLoadThreadMemory", data.threadId);

      instance.retrievalThreadId = data.threadId;

      // Update the text splitter
      await executeOnMain(
        NODE,
        "toolRetrievalUpdateTextSplitter",
        data.thread.assistants[0].tools[0]?.settings?.chunk_size ?? 4000,
        data.thread.assistants[0].tools[0]?.settings?.chunk_overlap ?? 200
      );
    }

    if (latestMessage.content) {
      const prompt =
        typeof latestMessage.content === "string"
          ? latestMessage.content
          : latestMessage.content[0].text;
      // Retrieve the result
      console.debug("toolRetrievalQuery", latestMessage.content);
      const retrievalResult = await executeOnMain(
        NODE,
        "toolRetrievalQueryResult",
        prompt
      );

      // Update the message content
      // Using the retrieval template with the result and query
      if (data.thread?.assistants[0].tools)
        data.messages[data.messages.length - 1].content =
          data.thread.assistants[0].tools[0].settings?.retrieval_template
            ?.replace("{CONTEXT}", retrievalResult)
            .replace("{QUESTION}", prompt);
    }

    // Filter out all the messages that are not text
    data.messages = data.messages.map((message) => {
      if (
        message.content &&
        typeof message.content !== "string" &&
        (message.content.length ?? 0) > 0
      ) {
        return {
          ...message,
          content: [message.content[0]],
        };
      }
      return message;
    });

    // Reroute the result to inference engine
    const output = {
      ...data,
      model: {
        ...data.model,
        engine: data.model.proxyEngine,
      },
    };
    events.emit(MessageEvent.OnMessageSent, output);
  }

  /**
   * Called when the extension is unloaded.
   */
  onUnload(): void {}

  async createAssistant(assistant: Assistant): Promise<void> {
    const assistantDir = await joinPath([
      JanAssistantExtension._homeDir,
      assistant.id,
    ]);
    if (!(await fs.existsSync(assistantDir))) await fs.mkdirSync(assistantDir);

    // store the assistant metadata json
    const assistantMetadataPath = await joinPath([
      assistantDir,
      "assistant.json",
    ]);
    try {
      await fs.writeFileSync(
        assistantMetadataPath,
        JSON.stringify(assistant, null, 2)
      );
    } catch (err) {
      console.error(err);
    }
  }

  async getAssistants(): Promise<Assistant[]> {
    // get all the assistant directories
    // get all the assistant metadata json
    const results: Assistant[] = [];
    const allFileName: string[] = await fs.readdirSync(
      JanAssistantExtension._homeDir
    );
    for (const fileName of allFileName) {
      const filePath = await joinPath([
        JanAssistantExtension._homeDir,
        fileName,
      ]);

      if (filePath.includes(".DS_Store")) continue;
      const jsonFiles: string[] = (await fs.readdirSync(filePath)).filter(
        (file: string) => file === "assistant.json"
      );

      if (jsonFiles.length !== 1) {
        // has more than one assistant file -> ignore
        continue;
      }

      const content = await fs.readFileSync(
        await joinPath([filePath, jsonFiles[0]]),
        "utf-8"
      );
      const assistant: Assistant =
        typeof content === "object" ? content : JSON.parse(content);

      results.push(assistant);
    }

    return results;
  }

  async deleteAssistant(assistant: Assistant): Promise<void> {
    if (assistant.id === "jan") {
      return Promise.reject("Cannot delete Jan Assistant");
    }

    // remove the directory
    const assistantDir = await joinPath([
      JanAssistantExtension._homeDir,
      assistant.id,
    ]);
    await fs.rmdirSync(assistantDir);
    return Promise.resolve();
  }

  private async createJanAssistant(): Promise<void> {
    const janAssistant: Assistant = {
      avatar: "",
      thread_location: undefined,
      id: "jan",
      object: "assistant",
      created_at: Date.now(),
      name: "Jan",
      description: "A default assistant that can use all downloaded models",
      model: "*",
      instructions: "",
      tools: [
        {
          type: "retrieval",
          enabled: false,
          settings: {
            top_k: 2,
            chunk_size: 1024,
            chunk_overlap: 64,
            retrieval_template: `Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
            ----------------
            CONTEXT: {CONTEXT}
            ----------------
            QUESTION: {QUESTION}
            ----------------
            Helpful Answer:`,
          },
        },
      ],
      file_ids: [],
      metadata: undefined,
    };

    await this.createAssistant(janAssistant);
  }
}