timmy-tower/artifacts/api-server/src/routes/gemini.ts

import { Router, type Request, type Response } from "express";
import { eq } from "drizzle-orm";
import { db, conversations, messages } from "@workspace/db";
import { ai, generateImage } from "@workspace/integrations-gemini-ai";
import { makeLogger } from "../lib/logger.js";

const router = Router();
const logger = makeLogger("gemini");

const DEFAULT_MODEL = "gemini-2.5-flash";

router.get("/conversations", async (_req: Request, res: Response) => {
  try {
    const rows = await db.select().from(conversations).orderBy(conversations.createdAt);
    res.json(rows);
  } catch (err) {
    logger.error("list conversations error", { error: err });
    res.status(500).json({ error: "Failed to list conversations" });
  }
});

router.post("/conversations", async (req: Request, res: Response) => {
  const { title } = req.body ?? {};
  if (typeof title !== "string" || !title.trim()) {
    res.status(400).json({ error: "title is required" });
    return;
  }
  try {
    const [row] = await db.insert(conversations).values({ title: title.trim() }).returning();
    res.status(201).json(row);
  } catch (err) {
    logger.error("create conversation error", { error: err });
    res.status(500).json({ error: "Failed to create conversation" });
  }
});

router.get("/conversations/:id", async (req: Request, res: Response) => {
  const id = parseInt(String(req.params["id"]), 10);
  if (isNaN(id)) {
    res.status(400).json({ error: "Invalid id" });
    return;
  }
  try {
    const [conv] = await db.select().from(conversations).where(eq(conversations.id, id));
    if (!conv) {
      res.status(404).json({ error: "Conversation not found" });
      return;
    }
    const msgs = await db
      .select()
      .from(messages)
      .where(eq(messages.conversationId, id))
      .orderBy(messages.createdAt);
    res.json({ ...conv, messages: msgs });
  } catch (err) {
    logger.error("get conversation error", { error: err });
    res.status(500).json({ error: "Failed to get conversation" });
  }
});

router.delete("/conversations/:id", async (req: Request, res: Response) => {
  const id = parseInt(String(req.params["id"]), 10);
  if (isNaN(id)) {
    res.status(400).json({ error: "Invalid id" });
    return;
  }
  try {
    const [conv] = await db.select().from(conversations).where(eq(conversations.id, id));
    if (!conv) {
      res.status(404).json({ error: "Conversation not found" });
      return;
    }
    await db.delete(conversations).where(eq(conversations.id, id));
    res.status(204).send();
  } catch (err) {
    logger.error("delete conversation error", { error: err });
    res.status(500).json({ error: "Failed to delete conversation" });
  }
});

router.get("/conversations/:id/messages", async (req: Request, res: Response) => {
  const id = parseInt(String(req.params["id"]), 10);
  if (isNaN(id)) {
    res.status(400).json({ error: "Invalid id" });
    return;
  }
  try {
    const msgs = await db
      .select()
      .from(messages)
      .where(eq(messages.conversationId, id))
      .orderBy(messages.createdAt);
    res.json(msgs);
  } catch (err) {
    logger.error("list messages error", { error: err });
    res.status(500).json({ error: "Failed to list messages" });
  }
});

router.post("/conversations/:id/messages", async (req: Request, res: Response) => {
  const conversationId = parseInt(String(req.params["id"]), 10);
  if (isNaN(conversationId)) {
    res.status(400).json({ error: "Invalid id" });
    return;
  }

  const { content, model } = req.body ?? {};
  if (typeof content !== "string" || !content.trim()) {
    res.status(400).json({ error: "content is required" });
    return;
  }

  try {
    const [conv] = await db.select().from(conversations).where(eq(conversations.id, conversationId));
    if (!conv) {
      res.status(404).json({ error: "Conversation not found" });
      return;
    }

    await db.insert(messages).values({ conversationId, role: "user", content: content.trim() });

    const history = await db
      .select()
      .from(messages)
      .where(eq(messages.conversationId, conversationId))
      .orderBy(messages.createdAt);

    const geminiContents = history.map((m) => ({
      role: m.role === "assistant" ? "model" : "user",
      parts: [{ text: m.content }],
    }));

    res.setHeader("Content-Type", "text/event-stream");
    res.setHeader("Cache-Control", "no-cache");
    res.setHeader("Connection", "keep-alive");
    res.flushHeaders();

    const sendEvent = (data: string) => {
      res.write(`data: ${data}\n\n`);
    };

    const stream = await ai.models.generateContentStream({
      model: model ?? DEFAULT_MODEL,
      contents: geminiContents,
    });

    let fullText = "";
    for await (const chunk of stream) {
      const text = chunk.text ?? "";
      if (text) {
        fullText += text;
        sendEvent(JSON.stringify({ text }));
      }
    }

    sendEvent(JSON.stringify({ done: true }));
    res.end();

    await db.insert(messages).values({ conversationId, role: "assistant", content: fullText });
  } catch (err) {
    logger.error("send message error", { error: err });
    if (!res.headersSent) {
      res.status(500).json({ error: "Failed to send message" });
    } else {
      res.write(`data: ${JSON.stringify({ error: "Stream error" })}\n\n`);
      res.end();
    }
  }
});

router.post("/generate-image", async (req: Request, res: Response) => {
  const { prompt } = req.body ?? {};
  if (typeof prompt !== "string" || !prompt.trim()) {
    res.status(400).json({ error: "prompt is required" });
    return;
  }
  try {
    const result = await generateImage(prompt.trim());
    res.json(result);
  } catch (err) {
    logger.error("generate image error", { error: err });
    res.status(500).json({ error: "Failed to generate image" });
  }
});

export default router;