ubiquitous-invention/apps/web/server/routers/ai.ts

import { createOpenAI } from "@ai-sdk/openai";
import { generateText } from "ai";
import { TRPCError } from "@trpc/server";
import { eq } from "drizzle-orm";
import { z } from "zod";
import { db as dbInstance } from "@tasks/database";
import { objects } from "@tasks/database/schema";
import { router, protectedProcedure } from "@/server/trpc";

type Db = typeof dbInstance;

const messageSchema = z.object({
  role: z.enum(["user", "assistant"]),
  content: z.string(),
});

const chatInputSchema = z.object({
  messages: z.array(messageSchema).min(1),
  context: z
    .object({
      workspaceId: z.string().optional(),
      objectId: z.string().uuid().optional(),
    })
    .optional(),
});

const suggestInputSchema = z.object({
  objectId: z.string().uuid().optional(),
  objectType: z.string().optional(),
});

const BASE_SYSTEM = `You are a helpful AI assistant embedded in a project management and collaboration app. Users organize work in workspaces with objects such as projects, tasks, documents, and groups. You help them plan work, clarify requirements, break down tasks, summarize content, and suggest next steps. Be concise, actionable, and friendly. Use markdown when it improves readability (bold, lists, short code snippets).`;

async function fetchObjectSummary(database: Db, objectId: string): Promise<string | null> {
  const row = await database.query.objects.findFirst({
    where: eq(objects.id, objectId),
    columns: {
      id: true,
      title: true,
      type: true,
      status: true,
      description: true,
      workspaceId: true,
    },
  });
  if (!row) return null;
  const bits = [
    `Object: ${row.title} (${row.type})`,
    row.status ? `Status: ${row.status}` : null,
    row.description ? `Description:\n${row.description}` : null,
    row.workspaceId ? `Workspace ID: ${row.workspaceId}` : null,
  ].filter(Boolean);
  return bits.join("\n");
}

function mockChatResponse(lastUser: string): string {
  return [
    "**Demo mode** — set `OPENAI_API_KEY` in your environment to use live AI.",
    "",
    "Here is a mock reply based on what you sent:",
    "",
    `> ${lastUser.slice(0, 280)}${lastUser.length > 280 ? "…" : ""}`,
    "",
    "In the real app, I would help you plan tasks, refine descriptions, and suggest next steps in your workspace.",
  ].join("\n");
}

async function callOpenAIChat(params: {
  system: string;
  messages: { role: "user" | "assistant"; content: string }[];
}): Promise<string> {
  const apiKey = process.env.OPENAI_API_KEY;
  if (!apiKey) {
    const lastUser = [...params.messages].reverse().find((m) => m.role === "user");
    return mockChatResponse(lastUser?.content ?? "");
  }

  try {
    const openai = createOpenAI({ apiKey });
    const { text } = await generateText({
      model: openai("gpt-4o"),
      system: params.system,
      messages: params.messages,
      temperature: 0.7,
      maxTokens: 4096,
    });
    const trimmed = text.trim();
    if (!trimmed) {
      throw new TRPCError({
        code: "INTERNAL_SERVER_ERROR",
        message: "Empty AI response",
      });
    }
    return trimmed;
  } catch (e) {
    if (e instanceof TRPCError) throw e;
    console.error("[ai.chat] generateText failed:", e);
    throw new TRPCError({
      code: "INTERNAL_SERVER_ERROR",
      message: "Failed to get AI response",
    });
  }
}

function suggestionsForContext(input: z.infer<typeof suggestInputSchema>): string[] {
  const t = (input.objectType ?? "").toLowerCase();
  const base = [
    "Create subtasks",
    "Set due date",
    "Add description",
    "Summarize open questions",
    "List blockers and next steps",
  ];
  if (t === "task" || t === "project") {
    return [
      "Break this into subtasks",
      "Suggest acceptance criteria",
      "Estimate effort",
      "Draft a checklist",
      ...base,
    ];
  }
  if (t === "document" || t === "whiteboard") {
    return [
      "Summarize this document",
      "Extract action items",
      "Outline key sections",
      "Generate task descriptions",
      ...base,
    ];
  }
  return ["Create a project plan", "Summarize this document", "Generate task descriptions", ...base];
}

export const aiRouter = router({
  chat: protectedProcedure.input(chatInputSchema).mutation(async ({ ctx, input }) => {
    let system = BASE_SYSTEM;
    const ctxParts: string[] = [];

    if (input.context?.workspaceId) {
      ctxParts.push(`Current workspace context ID: ${input.context.workspaceId}`);
    }

    if (input.context?.objectId) {
      const summary = await fetchObjectSummary(ctx.db, input.context.objectId);
      if (summary) {
        ctxParts.push("The user is focused on this object:\n" + summary);
      } else {
        ctxParts.push(
          `The user referenced object ID ${input.context.objectId}, but it was not found.`,
        );
      }
    }

    if (ctxParts.length) {
      system += "\n\n---\n" + ctxParts.join("\n\n");
    }

    const text = await callOpenAIChat({
      system,
      messages: input.messages,
    });

    return { text };
  }),

  suggestActions: protectedProcedure.input(suggestInputSchema).query(async ({ ctx, input }) => {
    let objectType = input.objectType;
    if (input.objectId && !objectType) {
      const row = await ctx.db.query.objects.findFirst({
        where: eq(objects.id, input.objectId),
        columns: { type: true },
      });
      objectType = row?.type;
    }
    return { actions: suggestionsForContext({ ...input, objectType }) };
  }),
});