path: root/apps/workers/openaiWorker.ts

import { and, Column, eq, inArray, sql } from "drizzle-orm";
import { DequeuedJob, Runner } from "liteque";
import { buildImpersonatingTRPCClient } from "trpc";
import { z } from "zod";

import type { InferenceClient } from "@karakeep/shared/inference";
import type { ZOpenAIRequest } from "@karakeep/shared/queues";
import { db } from "@karakeep/db";
import {
  bookmarks,
  bookmarkTags,
  customPrompts,
  tagsOnBookmarks,
} from "@karakeep/db/schema";
import { readAsset } from "@karakeep/shared/assetdb";
import serverConfig from "@karakeep/shared/config";
import { InferenceClientFactory } from "@karakeep/shared/inference";
import logger from "@karakeep/shared/logger";
import { buildImagePrompt, buildTextPrompt } from "@karakeep/shared/prompts";
import {
  OpenAIQueue,
  triggerRuleEngineOnEvent,
  triggerSearchReindex,
  triggerWebhook,
  zOpenAIRequestSchema,
} from "@karakeep/shared/queues";

const openAIResponseSchema = z.object({
  tags: z.array(z.string()),
});

function tagNormalizer(col: Column) {
  function normalizeTag(tag: string) {
    return tag.toLowerCase().replace(/[ \-_]/g, "");
  }

  return {
    normalizeTag,
    sql: sql`lower(replace(replace(replace(${col}, ' ', ''), '-', ''), '_', ''))`,
  };
}

async function attemptMarkTaggingStatus(
  jobData: object | undefined,
  status: "success" | "failure",
) {
  if (!jobData) {
    return;
  }
  try {
    const request = zOpenAIRequestSchema.parse(jobData);
    await db
      .update(bookmarks)
      .set({
        taggingStatus: status,
      })
      .where(eq(bookmarks.id, request.bookmarkId));
  } catch (e) {
    logger.error(`Something went wrong when marking the tagging status: ${e}`);
  }
}

export class OpenAiWorker {
  static build() {
    logger.info("Starting inference worker ...");
    const worker = new Runner<ZOpenAIRequest>(
      OpenAIQueue,
      {
        run: runOpenAI,
        onComplete: async (job) => {
          const jobId = job.id;
          logger.info(`[inference][${jobId}] Completed successfully`);
          await attemptMarkTaggingStatus(job.data, "success");
        },
        onError: async (job) => {
          const jobId = job.id;
          logger.error(
            `[inference][${jobId}] inference job failed: ${job.error}\n${job.error.stack}`,
          );
          if (job.numRetriesLeft == 0) {
            await attemptMarkTaggingStatus(job?.data, "failure");
          }
        },
      },
      {
        concurrency: 1,
        pollIntervalMs: 1000,
        timeoutSecs: serverConfig.inference.jobTimeoutSec,
      },
    );

    return worker;
  }
}

async function buildPrompt(
  bookmark: NonNullable<Awaited<ReturnType<typeof fetchBookmark>>>,
) {
  const prompts = await fetchCustomPrompts(bookmark.userId, "text");
  if (bookmark.link) {
    if (!bookmark.link.description && !bookmark.link.content) {
      throw new Error(
        `No content found for link "${bookmark.id}". Skipping ...`,
      );
    }

    const content = bookmark.link.content;
    return buildTextPrompt(
      serverConfig.inference.inferredTagLang,
      prompts,
      `URL: ${bookmark.link.url}
Title: ${bookmark.link.title ?? ""}
Description: ${bookmark.link.description ?? ""}
Content: ${content ?? ""}`,
      serverConfig.inference.contextLength,
    );
  }

  if (bookmark.text) {
    return buildTextPrompt(
      serverConfig.inference.inferredTagLang,
      prompts,
      bookmark.text.text ?? "",
      serverConfig.inference.contextLength,
    );
  }

  throw new Error("Unknown bookmark type");
}

async function fetchBookmark(linkId: string) {
  return await db.query.bookmarks.findFirst({
    where: eq(bookmarks.id, linkId),
    with: {
      link: true,
      text: true,
      asset: true,
    },
  });
}

async function inferTagsFromImage(
  jobId: string,
  bookmark: NonNullable<Awaited<ReturnType<typeof fetchBookmark>>>,
  inferenceClient: InferenceClient,
  abortSignal: AbortSignal,
) {
  const { asset, metadata } = await readAsset({
    userId: bookmark.userId,
    assetId: bookmark.asset.assetId,
  });

  if (!asset) {
    throw new Error(
      `[inference][${jobId}] AssetId ${bookmark.asset.assetId} for bookmark ${bookmark.id} not found`,
    );
  }

  const base64 = asset.toString("base64");
  return inferenceClient.inferFromImage(
    buildImagePrompt(
      serverConfig.inference.inferredTagLang,
      await fetchCustomPrompts(bookmark.userId, "images"),
    ),
    metadata.contentType,
    base64,
    { schema: openAIResponseSchema, abortSignal },
  );
}

async function fetchCustomPrompts(
  userId: string,
  appliesTo: "text" | "images",
) {
  const prompts = await db.query.customPrompts.findMany({
    where: and(
      eq(customPrompts.userId, userId),
      inArray(customPrompts.appliesTo, ["all_tagging", appliesTo]),
    ),
    columns: {
      text: true,
    },
  });

  let promptTexts = prompts.map((p) => p.text);
  if (containsTagsPlaceholder(prompts)) {
    promptTexts = await replaceTagsPlaceholders(promptTexts, userId);
  }

  return promptTexts;
}

async function replaceTagsPlaceholders(
  prompts: string[],
  userId: string,
): Promise<string[]> {
  const api = await buildImpersonatingTRPCClient(userId);
  const tags = (await api.tags.list()).tags;
  const tagsString = `[${tags.map((tag) => tag.name).join(", ")}]`;
  const aiTagsString = `[${tags
    .filter((tag) => tag.numBookmarksByAttachedType.human ?? 0 == 0)
    .map((tag) => tag.name)
    .join(", ")}]`;
  const userTagsString = `[${tags
    .filter((tag) => tag.numBookmarksByAttachedType.human ?? 0 > 0)
    .map((tag) => tag.name)
    .join(", ")}]`;

  return prompts.map((p) =>
    p
      .replaceAll("$tags", tagsString)
      .replaceAll("$aiTags", aiTagsString)
      .replaceAll("$userTags", userTagsString),
  );
}

function containsTagsPlaceholder(prompts: { text: string }[]): boolean {
  return (
    prompts.filter(
      (p) =>
        p.text.includes("$tags") ||
        p.text.includes("$aiTags") ||
        p.text.includes("$userTags"),
    ).length > 0
  );
}

async function inferTagsFromPDF(
  _jobId: string,
  bookmark: NonNullable<Awaited<ReturnType<typeof fetchBookmark>>>,
  inferenceClient: InferenceClient,
  abortSignal: AbortSignal,
) {
  const prompt = buildTextPrompt(
    serverConfig.inference.inferredTagLang,
    await fetchCustomPrompts(bookmark.userId, "text"),
    `Content: ${bookmark.asset.content}`,
    serverConfig.inference.contextLength,
  );
  return inferenceClient.inferFromText(prompt, {
    schema: openAIResponseSchema,
    abortSignal,
  });
}

async function inferTagsFromText(
  bookmark: NonNullable<Awaited<ReturnType<typeof fetchBookmark>>>,
  inferenceClient: InferenceClient,
  abortSignal: AbortSignal,
) {
  return await inferenceClient.inferFromText(await buildPrompt(bookmark), {
    schema: openAIResponseSchema,
    abortSignal,
  });
}

async function inferTags(
  jobId: string,
  bookmark: NonNullable<Awaited<ReturnType<typeof fetchBookmark>>>,
  inferenceClient: InferenceClient,
  abortSignal: AbortSignal,
) {
  let response;
  if (bookmark.link || bookmark.text) {
    response = await inferTagsFromText(bookmark, inferenceClient, abortSignal);
  } else if (bookmark.asset) {
    switch (bookmark.asset.assetType) {
      case "image":
        response = await inferTagsFromImage(
          jobId,
          bookmark,
          inferenceClient,
          abortSignal,
        );
        break;
      case "pdf":
        response = await inferTagsFromPDF(
          jobId,
          bookmark,
          inferenceClient,
          abortSignal,
        );
        break;
      default:
        throw new Error(`[inference][${jobId}] Unsupported bookmark type`);
    }
  } else {
    throw new Error(`[inference][${jobId}] Unsupported bookmark type`);
  }

  if (!response) {
    throw new Error(`[inference][${jobId}] Inference response is empty`);
  }

  try {
    let tags = openAIResponseSchema.parse(JSON.parse(response.response)).tags;
    logger.info(
      `[inference][${jobId}] Inferring tag for bookmark "${bookmark.id}" used ${response.totalTokens} tokens and inferred: ${tags}`,
    );

    // Sometimes the tags contain the hashtag symbol, let's strip them out if they do.
    // Additionally, trim the tags to prevent whitespaces at the beginning/the end of the tag.
    tags = tags.map((t) => {
      let tag = t;
      if (tag.startsWith("#")) {
        tag = t.slice(1);
      }
      return tag.trim();
    });

    return tags;
  } catch (e) {
    const responseSneak = response.response.substring(0, 20);
    throw new Error(
      `[inference][${jobId}] The model ignored our prompt and didn't respond with the expected JSON: ${JSON.stringify(e)}. Here's a sneak peak from the response: ${responseSneak}`,
    );
  }
}

async function connectTags(
  bookmarkId: string,
  inferredTags: string[],
  userId: string,
) {
  if (inferredTags.length == 0) {
    return;
  }

  await db.transaction(async (tx) => {
    // Attempt to match exiting tags with the new ones
    const { matchedTagIds, notFoundTagNames } = await (async () => {
      const { normalizeTag, sql: normalizedTagSql } = tagNormalizer(
        bookmarkTags.name,
      );
      const normalizedInferredTags = inferredTags.map((t) => ({
        originalTag: t,
        normalizedTag: normalizeTag(t),
      }));

      const matchedTags = await tx.query.bookmarkTags.findMany({
        where: and(
          eq(bookmarkTags.userId, userId),
          inArray(
            normalizedTagSql,
            normalizedInferredTags.map((t) => t.normalizedTag),
          ),
        ),
      });

      const matchedTagIds = matchedTags.map((r) => r.id);
      const notFoundTagNames = normalizedInferredTags
        .filter(
          (t) =>
            !matchedTags.some(
              (mt) => normalizeTag(mt.name) === t.normalizedTag,
            ),
        )
        .map((t) => t.originalTag);

      return { matchedTagIds, notFoundTagNames };
    })();

    // Create tags that didn't exist previously
    let newTagIds: string[] = [];
    if (notFoundTagNames.length > 0) {
      newTagIds = (
        await tx
          .insert(bookmarkTags)
          .values(
            notFoundTagNames.map((t) => ({
              name: t,
              userId,
            })),
          )
          .onConflictDoNothing()
          .returning()
      ).map((t) => t.id);
    }

    // Delete old AI tags
    const detachedTags = await tx
      .delete(tagsOnBookmarks)
      .where(
        and(
          eq(tagsOnBookmarks.attachedBy, "ai"),
          eq(tagsOnBookmarks.bookmarkId, bookmarkId),
        ),
      )
      .returning();

    const allTagIds = new Set([...matchedTagIds, ...newTagIds]);

    // Attach new ones
    const attachedTags = await tx
      .insert(tagsOnBookmarks)
      .values(
        [...allTagIds].map((tagId) => ({
          tagId,
          bookmarkId,
          attachedBy: "ai" as const,
        })),
      )
      .onConflictDoNothing()
      .returning();

    await triggerRuleEngineOnEvent(bookmarkId, [
      ...detachedTags.map((t) => ({
        type: "tagRemoved" as const,
        tagId: t.tagId,
      })),
      ...attachedTags.map((t) => ({
        type: "tagAdded" as const,
        tagId: t.tagId,
      })),
    ]);
  });
}

async function runOpenAI(job: DequeuedJob<ZOpenAIRequest>) {
  const jobId = job.id;

  const inferenceClient = InferenceClientFactory.build();
  if (!inferenceClient) {
    logger.debug(
      `[inference][${jobId}] No inference client configured, nothing to do now`,
    );
    return;
  }

  const request = zOpenAIRequestSchema.safeParse(job.data);
  if (!request.success) {
    throw new Error(
      `[inference][${jobId}] Got malformed job request: ${request.error.toString()}`,
    );
  }

  const { bookmarkId } = request.data;
  const bookmark = await fetchBookmark(bookmarkId);
  if (!bookmark) {
    throw new Error(
      `[inference][${jobId}] bookmark with id ${bookmarkId} was not found`,
    );
  }

  logger.info(
    `[inference][${jobId}] Starting an inference job for bookmark with id "${bookmark.id}"`,
  );

  const tags = await inferTags(
    jobId,
    bookmark,
    inferenceClient,
    job.abortSignal,
  );

  await connectTags(bookmarkId, tags, bookmark.userId);

  // Trigger a webhook
  await triggerWebhook(bookmarkId, "ai tagged");

  // Update the search index
  await triggerSearchReindex(bookmarkId);
}