SPB Git forge

spb/polyllm

Public
15commits 1branches 0releases
2.2 MBsize
maindefault branch
13 days agolast push
TypeScript 97.4% SQL 1% JavaScript 0.9% CSS 0.6%
9.7 KB · 235 lines typescript
Raw Blame History
1import "server-only";2import { and, desc, eq, ilike, inArray, isNull, sql } from "drizzle-orm";3import { getDb, projectFiles, projects, messageAttachments, type ProjectFile } from "@/db";4import { ids, newId } from "@/lib/ids";5import { ApiError } from "@/lib/api";6import { IMAGE_MIME, isTextLike } from "@/lib/ai/core/content";7import { estimateAttachmentTokens } from "@/lib/client/tokens";89export const LIBRARY_MAX_BYTES = 10 * 1024 * 1024;10export const LIBRARY_MAX_TEXT_BYTES = 2 * 1024 * 1024;11export const LIBRARY_MAX_FILES = 500;1213/** Metadata only — the base64 payload never leaves the server except through GET /api/library/files/[id]. */14export function toPublicFile(f: ProjectFile) {15  return {16    id: f.id,17    projectId: f.projectId,18    kind: f.kind,19    name: f.name,20    mimeType: f.mimeType,21    sizeBytes: f.sizeBytes,22    width: f.width,23    height: f.height,24    estimatedTokens: f.estimatedTokens,25    description: f.description,26    createdAt: f.createdAt.toISOString(),27  };28}29export type PublicProjectFile = ReturnType<typeof toPublicFile>;3031/** Same classification as /api/attachments so anything accepted there is accepted here (and vice versa). */32export function kindOf(mime: string, name: string): string | null {33  if (IMAGE_MIME.has(mime)) return "image";34  if (mime === "application/pdf" || /\.pdf$/i.test(name)) return "pdf";35  if (mime === "text/csv" || /\.csv$/i.test(name)) return "csv";36  if (mime === "application/json" || /\.json$/i.test(name)) return "json";37  if (isTextLike(mime, name)) return /\.(md|txt)$/i.test(name) ? "text" : "code";38  return null;39}4041function pngSize(buf: Buffer): { width: number; height: number } | null {42  if (buf.length > 24 && buf.toString("ascii", 1, 4) === "PNG") return { width: buf.readUInt32BE(16), height: buf.readUInt32BE(20) };43  return null;44}4546/** JPEG SOF marker scan — enough for the token heuristic; returns null when the stream is unusual. */47function jpegSize(buf: Buffer): { width: number; height: number } | null {48  if (buf.length < 4 || buf[0] !== 0xff || buf[1] !== 0xd8) return null;49  let i = 2;50  while (i + 9 < buf.length) {51    if (buf[i] !== 0xff) return null;52    const marker = buf[i + 1];53    const len = buf.readUInt16BE(i + 2);54    if ((marker >= 0xc0 && marker <= 0xc3) || (marker >= 0xc5 && marker <= 0xc7) || (marker >= 0xc9 && marker <= 0xcb) || (marker >= 0xcd && marker <= 0xcf)) {55      return { height: buf.readUInt16BE(i + 5), width: buf.readUInt16BE(i + 7) };56    }57    i += 2 + len;58  }59  return null;60}6162export function imageSize(buf: Buffer, mime: string): { width: number; height: number } | null {63  if (mime === "image/png") return pngSize(buf);64  if (mime === "image/jpeg") return jpegSize(buf);65  return null;66}6768export interface ListFilesFilter {69  /** A project id, `"none"` for the global library, or undefined for everything the user owns. */70  projectId?: string | "none";71  q?: string;72  kind?: string;73  limit?: number;74}7576export async function listFiles(userId: string, f: ListFilesFilter = {}) {77  const conds = [eq(projectFiles.userId, userId)];78  if (f.projectId === "none") conds.push(isNull(projectFiles.projectId));79  else if (f.projectId) conds.push(eq(projectFiles.projectId, f.projectId));80  if (f.q) conds.push(ilike(projectFiles.name, `%${f.q.replace(/[%_]/g, "\\$&")}%`));81  if (f.kind) conds.push(eq(projectFiles.kind, f.kind));82  const rows = await getDb()83    .select({84      id: projectFiles.id,85      projectId: projectFiles.projectId,86      kind: projectFiles.kind,87      name: projectFiles.name,88      mimeType: projectFiles.mimeType,89      sizeBytes: projectFiles.sizeBytes,90      width: projectFiles.width,91      height: projectFiles.height,92      estimatedTokens: projectFiles.estimatedTokens,93      description: projectFiles.description,94      createdAt: projectFiles.createdAt,95    })96    .from(projectFiles)97    .where(and(...conds))98    .orderBy(desc(projectFiles.createdAt))99    .limit(Math.min(f.limit ?? 200, LIBRARY_MAX_FILES));100  return rows.map((r) => ({ ...r, createdAt: r.createdAt.toISOString() }) as PublicProjectFile);101}102103export async function getFile(userId: string, id: string): Promise<ProjectFile> {104  const [row] = await getDb()105    .select()106    .from(projectFiles)107    .where(and(eq(projectFiles.id, id), eq(projectFiles.userId, userId)))108    .limit(1);109  if (!row) throw new ApiError(404, "File not found", "NOT_FOUND");110  return row;111}112113export interface SaveFileInput {114  file: File;115  projectId?: string | null;116  description?: string | null;117}118119/** Validates, classifies, measures and stores an uploaded file (mirrors /api/attachments limits). */120export async function saveFile(userId: string, input: SaveFileInput) {121  const { file } = input;122  const name = (file.name || "file").slice(0, 200);123  const mime = (file.type || "application/octet-stream").toLowerCase();124  const kind = kindOf(mime, name);125  if (!kind) throw new ApiError(415, `Unsupported file type (${mime || "unknown"}). Use images, PDF, text, code, CSV or JSON.`, "UNSUPPORTED_TYPE");126  if (file.size > LIBRARY_MAX_BYTES) throw new ApiError(413, "File too large (max 10 MB)", "PAYLOAD_TOO_LARGE");127  if (kind !== "image" && kind !== "pdf" && file.size > LIBRARY_MAX_TEXT_BYTES) throw new ApiError(413, "Text files are limited to 2 MB", "PAYLOAD_TOO_LARGE");128  const db = getDb();129  if (input.projectId) {130    const [p] = await db131      .select({ id: projects.id })132      .from(projects)133      .where(and(eq(projects.id, input.projectId), eq(projects.userId, userId)))134      .limit(1);135    if (!p) throw new ApiError(404, "Project not found", "NOT_FOUND");136  }137  const [{ n }] = await db138    .select({ n: sql<number>`count(*)::int` })139    .from(projectFiles)140    .where(eq(projectFiles.userId, userId));141  if (n >= LIBRARY_MAX_FILES) throw new ApiError(409, `Library is full (${LIBRARY_MAX_FILES} files). Delete some files first.`, "LIBRARY_FULL");142  const buf = Buffer.from(await file.arrayBuffer());143  const dims = kind === "image" ? imageSize(buf, mime) : null;144  const effectiveMime = kind === "image" ? mime : kind === "pdf" ? "application/pdf" : mime.startsWith("text/") || mime === "application/json" ? mime : "text/plain";145  const estimatedTokens = estimateAttachmentTokens({ kind, sizeBytes: buf.length, width: dims?.width ?? null, height: dims?.height ?? null });146  const [row] = await db147    .insert(projectFiles)148    .values({149      id: newId("pfl"),150      userId,151      projectId: input.projectId ?? null,152      kind,153      name,154      mimeType: effectiveMime,155      sizeBytes: buf.length,156      dataBase64: buf.toString("base64"),157      width: dims?.width ?? null,158      height: dims?.height ?? null,159      estimatedTokens,160      description: input.description?.trim().slice(0, 500) || null,161    })162    .returning();163  return toPublicFile(row);164}165166export async function updateFile(userId: string, id: string, patch: { name?: string; description?: string | null; projectId?: string | null }) {167  const set: Partial<typeof projectFiles.$inferInsert> = {};168  if (patch.name !== undefined) set.name = patch.name.trim().slice(0, 200) || "file";169  if (patch.description !== undefined) set.description = patch.description?.trim().slice(0, 500) || null;170  if (patch.projectId !== undefined) {171    if (patch.projectId) {172      const [p] = await getDb()173        .select({ id: projects.id })174        .from(projects)175        .where(and(eq(projects.id, patch.projectId), eq(projects.userId, userId)))176        .limit(1);177      if (!p) throw new ApiError(404, "Project not found", "NOT_FOUND");178    }179    set.projectId = patch.projectId;180  }181  if (!Object.keys(set).length) return toPublicFile(await getFile(userId, id));182  const [row] = await getDb()183    .update(projectFiles)184    .set(set)185    .where(and(eq(projectFiles.id, id), eq(projectFiles.userId, userId)))186    .returning();187  if (!row) throw new ApiError(404, "File not found", "NOT_FOUND");188  return toPublicFile(row);189}190191export async function deleteFile(userId: string, id: string) {192  const res = await getDb()193    .delete(projectFiles)194    .where(and(eq(projectFiles.id, id), eq(projectFiles.userId, userId)))195    .returning({ id: projectFiles.id });196  if (!res.length) throw new ApiError(404, "File not found", "NOT_FOUND");197}198199export interface PendingAttachment {200  id: string;201  kind: string;202  name: string;203  mimeType: string;204  sizeBytes: number;205  width: number | null;206  height: number | null;207}208209/**210 * Copies library files into `message_attachments` as *pending* attachments (messageId / conversationId null),211 * exactly like a fresh upload through POST /api/attachments. The chat service links them to the user message on send.212 */213export async function attachFiles(userId: string, fileIds: string[]): Promise<PendingAttachment[]> {214  const uniq = Array.from(new Set(fileIds));215  if (!uniq.length) return [];216  const db = getDb();217  const rows = await db218    .select()219    .from(projectFiles)220    .where(and(eq(projectFiles.userId, userId), inArray(projectFiles.id, uniq)));221  if (rows.length !== uniq.length) throw new ApiError(404, "Some files were not found", "NOT_FOUND");222  const byId = new Map(rows.map((r) => [r.id, r]));223  const values = uniq.map((fid) => {224    const f = byId.get(fid)!;225    return { id: ids.attachment(), userId, kind: f.kind, name: f.name, mimeType: f.mimeType, sizeBytes: f.sizeBytes, dataBase64: f.dataBase64, width: f.width, height: f.height };226  });227  const inserted = await db228    .insert(messageAttachments)229    .values(values)230    .returning({ id: messageAttachments.id, kind: messageAttachments.kind, name: messageAttachments.name, mimeType: messageAttachments.mimeType, sizeBytes: messageAttachments.sizeBytes, width: messageAttachments.width, height: messageAttachments.height });231  // Preserve the caller's order.232  const order = new Map(values.map((v, i) => [v.id, i]));233  return inserted.sort((a, b) => (order.get(a.id) ?? 0) - (order.get(b.id) ?? 0));234}235