aboutsummaryrefslogtreecommitdiffstats
path: root/workers/openai.ts
blob: cc23f700aa1a94c7a56bdba12933d13593656d0f (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
import prisma, { BookmarkedLink, BookmarkedLinkDetails } from "@remember/db";
import logger from "@remember/shared/logger";
import { ZOpenAIRequest, zOpenAIRequestSchema } from "@remember/shared/queues";
import { Job } from "bullmq";
import OpenAI from "openai";
import { z } from "zod";

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

let openai: OpenAI | undefined;

if (process.env.OPENAI_API_KEY && process.env.OPENAI_ENABLED) {
  openai = new OpenAI({
    apiKey: process.env["OPENAI_API_KEY"], // This is the default and can be omitted
  });
}

function buildPrompt(url: string, description: string) {
  return `
You are a bot who given an article, extracts relevant "hashtags" out of them.
You must respond in JSON with the key "tags" and the value is list of tags.
----
URL: ${url}
Description: ${description}
  `;
}

async function fetchLink(linkId: string) {
  return await prisma.bookmarkedLink.findUnique({
    where: {
      id: linkId,
    },
    include: {
      details: true,
    },
  });
}

async function inferTags(
  jobId: string,
  link: BookmarkedLink,
  linkDetails: BookmarkedLinkDetails | null,
  openai: OpenAI,
) {
  const linkDescription = linkDetails?.description;
  if (!linkDescription) {
    throw new Error(
      `[openai][${jobId}] No description found for link "${link.id}". Skipping ...`,
    );
  }

  const chatCompletion = await openai.chat.completions.create({
    messages: [
      { role: "system", content: buildPrompt(link.url, linkDescription) },
    ],
    model: "gpt-3.5-turbo-0125",
    response_format: { type: "json_object" },
  });

  let response = chatCompletion.choices[0].message.content;
  if (!response) {
    throw new Error(`[openai][${jobId}] Got no message content from OpenAI`);
  }

  try {
    const tags = openAIResponseSchema.parse(JSON.parse(response)).tags;
    logger.info(
      `[openai][${jobId}] Inferring tag for url "${link.url}" used ${chatCompletion.usage?.total_tokens} tokens and inferred: ${tags}`,
    );
    return tags;
  } catch (e) {
    throw new Error(
      `[openai][${jobId}] Failed to parse JSON response from OpenAI: ${e}`,
    );
  }
}

async function createTags(tags: string[], userId: string) {
  const existingTags = await prisma.bookmarkTags.findMany({
    select: {
      id: true,
      name: true,
    },
    where: {
      userId,
      name: {
        in: tags,
      },
    },
  });

  const existingTagSet = new Set<string>(existingTags.map((t) => t.name));

  let newTags = tags.filter((t) => !existingTagSet.has(t));

  // TODO: Prisma doesn't support createMany in Sqlite
  let newTagObjects = await Promise.all(
    newTags.map((t) => {
      return prisma.bookmarkTags.create({
        data: {
          name: t,
          userId: userId,
        },
      });
    }),
  );

  return existingTags.map((t) => t.id).concat(newTagObjects.map((t) => t.id));
}

async function connectTags(linkId: string, tagIds: string[]) {
  // TODO: Prisma doesn't support createMany in Sqlite
  await Promise.all(
    tagIds.map((tagId) => {
      return prisma.tagsOnLinks.create({
        data: {
          tagId,
          linkId,
        },
      });
    }),
  );
}

export default async function runOpenAI(job: Job<ZOpenAIRequest, void>) {
  const jobId = job.id || "unknown";

  if (!openai) {
    logger.debug(
      `[openai][${jobId}] OpenAI is not configured, nothing to do now`,
    );
    return;
  }

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

  const { linkId } = request.data;
  const link = await fetchLink(linkId);
  if (!link) {
    throw new Error(`[openai][${jobId}] link with id ${linkId} was not found`);
  }

  const tags = await inferTags(jobId, link, link.details, openai);

  const tagIds = await createTags(tags, link.userId);
  await connectTags(linkId, tagIds);
}