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

// Background detection — runs three Phase-0-validated detection prompts
// against the current model and returns structured findings.
//
// Sequential not parallel: small local models share a KV-cache budget; three
// concurrent calls can OOM. Each call is independent so total wall time
// roughly equals 3× per-call time, but stability is much higher.
//
// Post-validation strips hallucinated element ids — the same defensive step
// that rescued ~5% of phase-0 detection runs.

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

// ─── Output types ────────────────────────────────────────────────────────

export type FindingKind = "assumption" | "risk" | "inconsistency";

export interface Finding {
  kind: FindingKind;
  text: string;
  linkedElementIds: string[];
  confidence: number;
  severity?: "low" | "medium" | "high";
  validationCode?: string;
}

export interface DetectResult {
  findings: Finding[];
  inputTokens: number;
  outputTokens: number;
  durationMs: number;
  strippedRefs: number;
  droppedFindings: number;
  provider: string;
  model: string;
}

// ─── Per-pass JSON schemas (oneOf-by-shape; small models cope) ──────────

const baseProps = {
  text: { type: "string", minLength: 1 },
  linkedElementIds: { type: "array", items: { type: "string" } },
  confidence: { type: "number", minimum: 0, maximum: 1 },
};

const assumptionsJsonSchema = {
  type: "object",
  additionalProperties: false,
  required: ["findings"],
  properties: {
    findings: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["text", "confidence"],
        properties: baseProps,
      },
    },
  },
} as const;

const risksJsonSchema = {
  type: "object",
  additionalProperties: false,
  required: ["findings"],
  properties: {
    findings: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["text", "confidence", "severity"],
        properties: {
          ...baseProps,
          severity: { type: "string", enum: ["low", "medium", "high"] },
        },
      },
    },
  },
} as const;

const inconsistenciesJsonSchema = {
  type: "object",
  additionalProperties: false,
  required: ["findings"],
  properties: {
    findings: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["text", "confidence"],
        properties: {
          ...baseProps,
          validationCode: { type: "string" },
        },
      },
    },
  },
} as const;

interface RawFinding {
  text?: string;
  linkedElementIds?: string[];
  confidence?: number;
  severity?: "low" | "medium" | "high";
  validationCode?: string;
}

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

export async function detectFindings(model: SysMLModel): Promise<DetectResult> {
  const gateway = defaultGateway();
  const start = Date.now();
  let totalIn = 0;
  let totalOut = 0;

  const userPayload = buildUserPayload(model);

  const assumptions = await runPass(gateway, "detect-assumptions", userPayload, assumptionsJsonSchema);
  totalIn += assumptions.inputTokens;
  totalOut += assumptions.outputTokens;

  const risks = await runPass(gateway, "detect-risks", userPayload, risksJsonSchema);
  totalIn += risks.inputTokens;
  totalOut += risks.outputTokens;

  const incons = await runPass(gateway, "detect-inconsistencies", userPayload, inconsistenciesJsonSchema);
  totalIn += incons.inputTokens;
  totalOut += incons.outputTokens;

  const raw: Finding[] = [
    ...assumptions.findings.map((f): Finding => ({
      kind: "assumption",
      text: f.text ?? "",
      linkedElementIds: f.linkedElementIds ?? [],
      confidence: clamp01(f.confidence ?? 0.5),
    })),
    ...risks.findings.map((f): Finding => ({
      kind: "risk",
      text: f.text ?? "",
      linkedElementIds: f.linkedElementIds ?? [],
      confidence: clamp01(f.confidence ?? 0.5),
      severity: f.severity ?? "medium",
    })),
    ...incons.findings.map((f): Finding => ({
      kind: "inconsistency",
      text: f.text ?? "",
      linkedElementIds: f.linkedElementIds ?? [],
      confidence: clamp01(f.confidence ?? 0.6),
      validationCode: f.validationCode,
    })),
  ];

  const validated = postValidate(raw, model);

  return {
    findings: validated.findings,
    inputTokens: totalIn,
    outputTokens: totalOut,
    durationMs: Date.now() - start,
    strippedRefs: validated.strippedRefs,
    droppedFindings: validated.droppedFindings,
    provider: gateway.provider,
    model: gateway.model,
  };
}

// ─── Per-pass runner ─────────────────────────────────────────────────────

interface PassResult {
  findings: RawFinding[];
  inputTokens: number;
  outputTokens: number;
}

async function runPass(
  gateway: ReturnType<typeof defaultGateway>,
  promptName: "detect-assumptions" | "detect-risks" | "detect-inconsistencies",
  userPayload: string,
  jsonSchema: object
): Promise<PassResult> {
  const prompt = loadPrompt(`socrates/${promptName}.md`);
  const messages: Message[] = [
    { role: "system", content: prompt },
    { role: "user", content: userPayload },
  ];

  try {
    const { value, result } = await chatJSON<{ findings?: RawFinding[] }>(gateway, messages, {
      temperature: 0.2,
      maxTokens: 768,
      jsonSchema: { name: "findings", schema: jsonSchema as Record<string, unknown> },
      jsonObjectMode: true,
      maxRepairs: 2,
    });
    return {
      findings: Array.isArray(value?.findings) ? value.findings : [],
      inputTokens: result.inputTokens,
      outputTokens: result.outputTokens,
    };
  } catch (err) {
    // Fail-soft: a busted detection pass shouldn't take down the whole detect.
    console.error(`[detect] ${promptName} pass failed:`, (err as Error).message);
    return { findings: [], inputTokens: 0, outputTokens: 0 };
  }
}

// ─── Trimmed payload (small models need tight context) ───────────────────

function buildUserPayload(model: SysMLModel): string {
  const compact = {
    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),
    })),
  };
  return `Model:\n\`\`\`json\n${JSON.stringify(compact, null, 2)}\n\`\`\`\n`;
}

// ─── Post-validation: strip hallucinated element refs ────────────────────

interface ValidateOut {
  findings: Finding[];
  strippedRefs: number;
  droppedFindings: number;
}

function postValidate(findings: Finding[], model: SysMLModel): ValidateOut {
  const validIds = collectValidIds(model);
  let stripped = 0;
  let dropped = 0;
  const out: Finding[] = [];

  for (const f of findings) {
    if (!f.text || f.text.trim().length === 0) {
      dropped++;
      continue;
    }
    const goodRefs: string[] = [];
    for (const ref of f.linkedElementIds) {
      if (validIds.has(ref)) goodRefs.push(ref);
      else stripped++;
    }
    // Drop only if it had refs and ALL of them were hallucinated.
    if (f.linkedElementIds.length > 0 && goodRefs.length === 0) {
      dropped++;
      continue;
    }
    out.push({ ...f, linkedElementIds: goodRefs });
  }
  return { findings: out, strippedRefs: stripped, droppedFindings: dropped };
}

function collectValidIds(model: SysMLModel): Set<string> {
  const ids = new Set<string>();
  for (const b of model.blocks) {
    ids.add(b.id);
    for (const p of b.properties) {
      ids.add(`${b.id}.${p.id}`);
      ids.add(p.name); // tolerate bare property names — common LLM pattern
    }
  }
  for (const a of model.associations) ids.add(a.id);
  for (const c of model.constraints) ids.add(c.id);
  for (const r of model.requirements) {
    ids.add(r.id);
    ids.add(r.tag);
  }
  return ids;
}

function clamp01(n: number): number {
  if (Number.isNaN(n)) return 0;
  return Math.max(0, Math.min(1, n));
}