Socrates/apps/web/lib/llm/socrates.ts

// Socrates conversation runner — server-side.
//
// Builds the LLM prompt (character + review + project context), wraps the
// thread history, asks for a structured JSON turn (`text` + optional
// `options`), persists user + assistant messages, and returns the new turn.

import "server-only";
import { defaultGateway, chatJSON, type Message } from "./gateway";
import { loadPrompt } from "./prompts";
import { prisma } from "../db/client";
import type { SysMLModel } from "../sysml/model";
import type { ValidationIssue } from "../sysml/validate";

// ─── Output schema for one Socrates turn ────────────────────────────────

export interface SocratesOption {
  n: number;
  label: string;
  sub?: string;
}

export interface SocratesTurn {
  text: string;
  options?: SocratesOption[];
}

const turnJsonSchema = {
  type: "object",
  additionalProperties: false,
  required: ["text"],
  properties: {
    text: { type: "string", minLength: 1 },
    options: {
      type: "array",
      maxItems: 3,
      items: {
        type: "object",
        additionalProperties: false,
        required: ["n", "label"],
        properties: {
          n: { type: "integer", minimum: 1, maximum: 3 },
          label: { type: "string", minLength: 1, maxLength: 60 },
          sub: { type: "string", maxLength: 80 },
        },
      },
    },
  },
} as const;

// ─── Public: handle one user turn → Socrates reply, persist both ────────

export interface SendUserTurnArgs {
  threadId: string;
  /** Caller passes the latest model so we don't re-load it inside this fn. */
  model: SysMLModel;
  issues?: ValidationIssue[];
  /** Body of the user's message. Empty string means "open the conversation". */
  userText: string;
}

export interface SendUserTurnResult {
  user: { id: string; createdAt: Date };
  assistant: {
    id: string;
    createdAt: Date;
    turn: SocratesTurn;
    provider: string;
    model: string;
    inputTokens: number;
    outputTokens: number;
  };
}

export async function sendUserTurn(args: SendUserTurnArgs): Promise<SendUserTurnResult> {
  const thread = await prisma.socratesThread.findUnique({
    where: { id: args.threadId },
    include: {
      messages: { orderBy: { createdAt: "asc" } },
    },
  });
  if (!thread) throw new Error(`thread ${args.threadId} not found`);

  // 1. Persist the user turn first so concurrent reads see it.
  let userRecord: { id: string; createdAt: Date } | undefined;
  if (args.userText.trim().length > 0) {
    userRecord = await prisma.socratesMessage.create({
      data: {
        threadId: thread.id,
        role: "user",
        content: JSON.stringify({ text: args.userText }),
      },
      select: { id: true, createdAt: true },
    });
  }

  // 2. Build the LLM context.
  const character = loadPrompt("socrates/character.md");
  const review = loadPrompt("socrates/review.md");

  const projectContext = JSON.stringify(
    {
      model: trimModel(args.model),
      issues: (args.issues ?? []).slice(0, 30).map(i => ({
        code: i.code,
        severity: i.severity,
        message: i.message,
        anchor: i.anchor,
      })),
    },
    null,
    2
  );

  const systemPrompt = [
    character,
    "---",
    review,
    "---",
    "Current project context (model + active validation issues):",
    "```json",
    projectContext,
    "```",
  ].join("\n\n");

  const historyMessages: Message[] = thread.messages.map(m => {
    const parsed = JSON.parse(m.content) as { text: string };
    return {
      role: m.role === "user" ? "user" : "assistant",
      content: parsed.text,
    };
  });

  if (args.userText.trim().length > 0) {
    historyMessages.push({ role: "user", content: args.userText });
  } else if (historyMessages.length === 0) {
    // Opening turn — give Socrates a kick.
    historyMessages.push({
      role: "user",
      content: "Open the conversation. Surface the most important tension you see in this model.",
    });
  }

  // 3. Ask the model.
  const gateway = defaultGateway();
  const { value, result } = await chatJSON<SocratesTurn>(
    gateway,
    [{ role: "system", content: systemPrompt }, ...historyMessages],
    {
      temperature: 0.4,
      maxTokens: 768,
      jsonSchema: { name: "socrates_turn", schema: turnJsonSchema as Record<string, unknown> },
      jsonObjectMode: true, // safety net for adapters lacking strict json_schema
      maxRepairs: 2,
    }
  );

  // Defensive normalization
  const turn: SocratesTurn = {
    text: typeof value?.text === "string" ? value.text : "[empty]",
    ...(Array.isArray(value?.options) && value.options.length > 0
      ? { options: value.options.slice(0, 3) }
      : {}),
  };

  // 4. Persist the assistant turn.
  const assistantRecord = await prisma.socratesMessage.create({
    data: {
      threadId: thread.id,
      role: "assistant",
      content: JSON.stringify(turn),
      provider: gateway.provider,
      model: gateway.model,
      inputTokens: result.inputTokens,
      outputTokens: result.outputTokens,
    },
    select: { id: true, createdAt: true },
  });

  await prisma.socratesThread.update({
    where: { id: thread.id },
    data: { updatedAt: new Date() },
  });

  return {
    user: userRecord ?? { id: "(skipped)", createdAt: new Date(0) },
    assistant: {
      id: assistantRecord.id,
      createdAt: assistantRecord.createdAt,
      turn,
      provider: gateway.provider,
      model: gateway.model,
      inputTokens: result.inputTokens,
      outputTokens: result.outputTokens,
    },
  };
}

// ─── Anchored threads (contextual mini-Socrates per finding) ─────────────

/** Find an existing thread anchored to `anchor` for this project, or create
 *  one. Returns the thread id. */
export async function getOrCreateAnchoredThread(
  projectId: string,
  anchor: string,
  title?: string
): Promise<string> {
  const existing = await prisma.socratesThread.findFirst({
    where: { projectId, anchorElementId: anchor },
    orderBy: { updatedAt: "desc" },
  });
  if (existing) return existing.id;
  const created = await prisma.socratesThread.create({
    data: { projectId, anchorElementId: anchor, status: "open", title: title ?? null },
  });
  return created.id;
}

export async function listAnchoredThreadMessages(threadId: string) {
  const rows = await prisma.socratesMessage.findMany({
    where: { threadId },
    orderBy: { createdAt: "asc" },
    select: { id: true, role: true, content: true, createdAt: true },
  });
  return rows.map(r => {
    let text = "";
    let options: SocratesOption[] | undefined;
    try {
      const parsed = JSON.parse(r.content) as { text?: string; options?: SocratesOption[] };
      text = parsed.text ?? "";
      options = parsed.options;
    } catch {
      text = r.content;
    }
    return { id: r.id, role: r.role, text, options, createdAt: r.createdAt };
  });
}

// ─── Helpers ─────────────────────────────────────────────────────────────

function trimModel(model: SysMLModel): unknown {
  // Compact view for Socrates — drop ids of properties (just names) and
  // multiplicity / type detail to save tokens.
  return {
    blocks: model.blocks.map(b => ({
      id: b.id,
      label: b.label,
      kind: b.kind,
      properties: b.properties.map(p => p.name),
    })),
    associations: model.associations.map(a => ({
      id: a.id,
      from: a.fromBlockId,
      to: a.toBlockId,
      label: a.label,
      kind: a.kind,
    })),
    constraints: model.constraints.map(c => ({
      id: c.id,
      label: c.label,
      appliesTo: c.appliesTo,
    })),
    requirements: model.requirements.map(r => ({
      id: r.id,
      tag: r.tag,
      text: r.text,
      satisfiedBy: r.relations
        .filter(rel => rel.kind === "satisfy")
        .map(rel => (rel as { kind: "satisfy"; blockId: string }).blockId),
    })),
  };
}