aegida-console / app / api / chat / route.ts
route.ts
Raw
import { AuthError, authenticateRequest } from "@/lib/auth/server";
import {
  createConversation,
  createMessage,
  getRetryPrompt,
  listGatewayMessages,
  ownedConversationExists,
  updateConversationModel,
  updateMessage,
} from "@/lib/db/conversations";
import { getPool } from "@/lib/db/pool";
import { attachPendingAttachments } from "@/lib/db/attachments";
import { GatewayError, streamGatewayResponse } from "@/lib/chat/gateway";
import { ChatValidationError, validatePersistentChatRequest } from "@/lib/chat/validation";

const noStoreHeaders = {
  "Cache-Control": "no-store",
  "X-Content-Type-Options": "nosniff",
};
const MAX_TITLE_LENGTH = 48;

export async function POST(request: Request): Promise<Response> {
  try {
    const user = await authenticateRequest(request);
    const authorization = request.headers.get("Authorization");
    if (!authorization) throw new AuthError();
    const payload = validatePersistentChatRequest(await request.json());
    const database = getPool();
    let conversationId = payload.conversationId;
    let userMessageId = "";

    if (conversationId) {
      if (!(await ownedConversationExists(database, user.id, conversationId))) {
        return errorResponse(404, "Чат не найден");
      }
      await updateConversationModel(database, user.id, conversationId, payload.model);
    } else {
      conversationId = await createConversation(
        database,
        user.id,
        payload.model,
        createTitle(payload.content ?? "Новый чат"),
      );
    }

    if (payload.retry) {
      const retryPrompt = await getRetryPrompt(database, user.id, conversationId);
      if (!retryPrompt) return errorResponse(400, "Некорректный запрос");
    } else {
      userMessageId = await createMessage(
        database,
        conversationId,
        "user",
        payload.content ?? "",
        "complete",
      );
      if (payload.attachmentIds?.length) {
        await attachPendingAttachments(
          database,
          user.id,
          userMessageId,
          payload.attachmentIds,
        );
      }
    }

    const gatewayMessages = await listGatewayMessages(database, user.id, conversationId);
    const assistantMessageId = await createMessage(
      database,
      conversationId,
      "assistant",
      "",
      "streaming",
    );
    let source: ReadableStream<Uint8Array>;
    try {
      source = await streamGatewayResponse(
        {
          model: payload.model,
          messages: gatewayMessages,
        },
        authorization,
        request.signal,
      );
    } catch (error) {
      await updateMessage(database, assistantMessageId, "", "error");
      if (error instanceof GatewayError) {
        return errorResponse(error.status, error.message, {
          "X-Conversation-Id": conversationId,
          "X-Assistant-Message-Id": assistantMessageId,
          "X-User-Message-Id": userMessageId,
        });
      }
      throw error;
    }
    const stream = persistStream(source, assistantMessageId);

    return new Response(stream, {
      headers: {
        ...noStoreHeaders,
        "Content-Type": "text/plain; charset=utf-8",
        "X-Conversation-Id": conversationId,
        "X-Assistant-Message-Id": assistantMessageId,
        "X-User-Message-Id": userMessageId,
      },
    });
  } catch (error) {
    return chatErrorResponse(error);
  }
}

function persistStream(source: ReadableStream<Uint8Array>, messageId: string) {
  const transform = new TransformStream<Uint8Array, Uint8Array>();
  const decoder = new TextDecoder();
  let content = "";

  void source
    .pipeThrough(
      new TransformStream({
        transform(chunk, controller) {
          content += decoder.decode(chunk, { stream: true });
          controller.enqueue(chunk);
        },
        flush() {
          content += decoder.decode();
        },
      }),
    )
    .pipeTo(transform.writable)
    .then(() => updateMessage(getPool(), messageId, content, "complete"))
    .catch(() => updateMessage(getPool(), messageId, content, "stopped"));

  return transform.readable;
}

export function chatErrorResponse(error: unknown): Response {
  if (error instanceof AuthError) return errorResponse(error.status, error.message);
  if (error instanceof ChatValidationError || error instanceof SyntaxError) {
    return errorResponse(400, "Некорректный запрос");
  }
  if (error instanceof GatewayError) return errorResponse(error.status, error.message);
  return errorResponse(503, "Хранилище чатов временно недоступно");
}

function errorResponse(
  status: number,
  message: string,
  headers: HeadersInit = {},
): Response {
  return Response.json(
    { error: message },
    { status, headers: { ...noStoreHeaders, ...Object.fromEntries(new Headers(headers)) } },
  );
}

function createTitle(content: string): string {
  const compact = content.replace(/\s+/g, " ").trim();
  return compact.length <= MAX_TITLE_LENGTH
    ? compact
    : `${compact.slice(0, MAX_TITLE_LENGTH - 1).trimEnd()}`;
}