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

// Seed → SysMLModel via the Phase-0-validated generate prompt.
//
// Returns a validated, expanded SysMLModel with property ids, multiplicities,
// stereotypes filled in. The LLM emits a "lean" shape; we expand it to match
// the canonical lib/sysml/model.ts types.

import "server-only";
import { defaultGateway, chatJSON, type Message } from "./gateway";
import { loadPrompt } from "./prompts";
import type {
  SysMLModel,
  Block,
  Association,
  Constraint,
  Requirement,
  Property,
  PropertyType,
} from "../sysml/model";
import type { SeedDraft } from "./seedInterview";

// ─── LLM-facing lean shape (matches Phase 0's generate.md output) ───────

interface LeanProperty {
  name: string;
  type: { kind: "string" | "number" | "boolean" | "enum"; values?: string[] };
}
interface LeanBlock {
  id: string;
  label: string;
  kind: "system" | "actor" | "block";
  properties?: LeanProperty[];
  confidence: number;
}
interface LeanAssociation {
  id: string;
  fromBlockId: string;
  toBlockId: string;
  label?: string;
  kind: "association" | "composition" | "aggregation" | "generalization" | "constraintApplies";
  confidence: number;
}
interface LeanConstraint {
  id: string;
  label: string;
  expression?: string;
  appliesTo?: string[];
  confidence: number;
}
interface LeanRequirement {
  id: string;
  tag: string;
  text: string;
  satisfiedBy?: string[];
  confidence: number;
}
interface LeanModel {
  systemOfInterestId?: string;
  blocks: LeanBlock[];
  associations?: LeanAssociation[];
  constraints?: LeanConstraint[];
  requirements?: LeanRequirement[];
  overallConfidence?: number;
  notes?: string;
}

// ─── JSON Schema for response_format ─────────────────────────────────────

const generateJsonSchema = {
  type: "object",
  additionalProperties: false,
  properties: {
    systemOfInterestId: { type: "string" },
    blocks: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "label", "kind", "confidence"],
        properties: {
          id: { type: "string", minLength: 1 },
          label: { type: "string", minLength: 1 },
          kind: { type: "string", enum: ["system", "actor", "block"] },
          confidence: { type: "number", minimum: 0, maximum: 1 },
          properties: {
            type: "array",
            items: {
              type: "object",
              additionalProperties: false,
              required: ["name", "type"],
              properties: {
                name: { type: "string", minLength: 1 },
                type: {
                  type: "object",
                  additionalProperties: false,
                  required: ["kind"],
                  properties: {
                    kind: { type: "string", enum: ["string", "number", "boolean", "enum"] },
                    values: { type: "array", items: { type: "string" } },
                  },
                },
              },
            },
          },
        },
      },
    },
    associations: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "fromBlockId", "toBlockId", "kind", "confidence"],
        properties: {
          id: { type: "string", minLength: 1 },
          fromBlockId: { type: "string", minLength: 1 },
          toBlockId: { type: "string", minLength: 1 },
          label: { type: "string" },
          kind: { type: "string", enum: ["association", "composition", "aggregation", "generalization", "constraintApplies"] },
          confidence: { type: "number", minimum: 0, maximum: 1 },
        },
      },
    },
    constraints: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "label", "confidence"],
        properties: {
          id: { type: "string", minLength: 1 },
          label: { type: "string", minLength: 1 },
          expression: { type: "string" },
          appliesTo: { type: "array", items: { type: "string" } },
          confidence: { type: "number", minimum: 0, maximum: 1 },
        },
      },
    },
    requirements: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "tag", "text", "confidence"],
        properties: {
          id: { type: "string", minLength: 1 },
          tag: { type: "string", minLength: 1 },
          text: { type: "string", minLength: 1 },
          satisfiedBy: { type: "array", items: { type: "string" } },
          confidence: { type: "number", minimum: 0, maximum: 1 },
        },
      },
    },
    overallConfidence: { type: "number", minimum: 0, maximum: 1 },
    notes: { type: "string" },
  },
  required: ["blocks"],
} as const;

// ─── Public API ──────────────────────────────────────────────────────────

export interface GenerateModelResult {
  model: SysMLModel;
  overallConfidence: number;
  notes?: string;
  inputTokens: number;
  outputTokens: number;
  provider: string;
  model_name: string;
}

export async function generateModelFromSeed(seed: SeedDraft): Promise<GenerateModelResult> {
  const character = loadPrompt("socrates/character.md");
  const generate = loadPrompt("socrates/generate.md");

  const userPayload =
    `Seed payload to model:\n\`\`\`json\n${JSON.stringify(seed, null, 2)}\n\`\`\`\n` +
    `Return only the JSON object conforming to the schema. No prose preamble, no code fences.`;

  const messages: Message[] = [
    { role: "system", content: `${character}\n\n---\n\n${generate}` },
    { role: "user", content: userPayload },
  ];

  const gateway = defaultGateway();
  const { value, result } = await chatJSON<LeanModel>(gateway, messages, {
    temperature: 0.3,
    maxTokens: 2048,
    jsonSchema: { name: "sysml_model", schema: generateJsonSchema as Record<string, unknown> },
    jsonObjectMode: true,
    maxRepairs: 2,
  });

  const model = expand(value);

  return {
    model,
    overallConfidence: clamp01(value?.overallConfidence ?? 0.5),
    notes: value?.notes,
    inputTokens: result.inputTokens,
    outputTokens: result.outputTokens,
    provider: gateway.provider,
    model_name: gateway.model,
  };
}

// ─── Lean → canonical SysMLModel ────────────────────────────────────────

function expand(lean: LeanModel | null | undefined): SysMLModel {
  if (!lean || !Array.isArray(lean.blocks)) {
    return { blocks: [], associations: [], constraints: [], requirements: [] };
  }

  const blocks: Block[] = lean.blocks.map(b => ({
    id: b.id,
    label: b.label,
    kind: b.kind,
    stereotypes: [b.kind],
    properties: (b.properties ?? []).map((p, i) => ({
      id: `${b.id}_p${i + 1}`,
      name: p.name,
      type: expandPropertyType(p.type),
      multiplicity: "0..1" as const,
    })),
  }));

  const associations: Association[] = (lean.associations ?? []).map(a => ({
    id: a.id,
    fromBlockId: a.fromBlockId,
    toBlockId: a.toBlockId,
    label: a.label ?? "",
    kind: a.kind,
  }));

  const constraints: Constraint[] = (lean.constraints ?? []).map(c => ({
    id: c.id,
    label: c.label,
    expression: c.expression ?? "",
    appliesTo: c.appliesTo ?? [],
  }));

  const requirements: Requirement[] = (lean.requirements ?? []).map(r => ({
    id: r.id,
    tag: r.tag,
    text: r.text,
    relations: (r.satisfiedBy ?? []).map(blockId => ({ kind: "satisfy" as const, blockId })),
  }));

  // Resolve SoI: prefer the LLM-supplied id if it exists; otherwise the unique kind:'system' block.
  let soiId = lean.systemOfInterestId;
  if (!soiId) {
    const systems = blocks.filter(b => b.kind === "system");
    if (systems.length === 1) soiId = systems[0]!.id;
  }

  return {
    systemOfInterestId: soiId,
    blocks,
    associations,
    constraints,
    requirements,
  };
}

function expandPropertyType(type: { kind: string; values?: string[] }): PropertyType {
  if (type.kind === "enum") return { kind: "enum", values: type.values ?? [] };
  if (type.kind === "number") return { kind: "number" };
  if (type.kind === "boolean") return { kind: "boolean" };
  return { kind: "string" };
}
function clamp01(n: number): number {
  if (Number.isNaN(n)) return 0;
  return Math.max(0, Math.min(1, n));
}
// Suppress unused-import warning in some TS configs.
type _Property = Property;