webapp/lib/api.ts

import type {
  ParseResult,
  BulkImportPreview,
  TrajectoryTree,
  TrajectoryTreeNode,
  TreeSummary,
  ValidationResult
} from "@/lib/tree/types";

async function requestJson<T>(url: string, init?: RequestInit) {
  const response = await fetch(url, init);
  const body = await response.json().catch(() => ({ error: response.statusText }));
  if (!response.ok) {
    throw new Error(body.error || response.statusText);
  }
  return body as T;
}

export function fetchTrees() {
  return requestJson<{ trees: TreeSummary[] }>("/api/trees");
}

export function createTreeFromRaw(input: { title?: string; rawText: string; nodes?: TrajectoryTreeNode[] }) {
  return requestJson<{ tree: TrajectoryTree; validation: ValidationResult }>("/api/trees", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify(input)
  });
}

export function previewBulkTreeImport(input: {
  rawJsonl: string;
  trajectoryColumn?: string;
  titleColumn?: string;
}) {
  return requestJson<{ preview: BulkImportPreview }>("/api/trees/bulk", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify(input)
  });
}

export function createTreesFromJsonl(input: {
  rawJsonl: string;
  trajectoryColumn: string;
  titleColumn?: string;
}) {
  return requestJson<{ preview: BulkImportPreview; tree: TrajectoryTree | null; trees: TreeSummary[] }>("/api/trees/bulk", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ ...input, create: true })
  });
}

export function fetchTree(treeId: string) {
  return requestJson<{ tree: TrajectoryTree; validation: ValidationResult }>(
    `/api/trees/${encodeURIComponent(treeId)}`
  );
}

export function updateTree(treeId: string, patch: Partial<TrajectoryTree>) {
  return requestJson<{ tree: TrajectoryTree; validation: ValidationResult }>(
    `/api/trees/${encodeURIComponent(treeId)}`,
    {
      method: "PATCH",
      headers: { "content-type": "application/json" },
      body: JSON.stringify(patch)
    }
  );
}

export function deleteTree(treeId: string) {
  return requestJson<{ ok: true }>(`/api/trees/${encodeURIComponent(treeId)}`, {
    method: "DELETE"
  });
}

export function deleteTrees(input: { ids?: string[]; all?: boolean }) {
  return requestJson<{ ok: true; deletedIds: string[] }>("/api/trees", {
    method: "DELETE",
    headers: { "content-type": "application/json" },
    body: JSON.stringify(input)
  });
}

export function parseTrajectory(rawText: string) {
  return requestJson<{ parse: ParseResult }>("/api/trees/parse", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ rawText })
  });
}

export function validateTreePayload(input: { id?: string; tree?: TrajectoryTree }) {
  return requestJson<{ validation: ValidationResult }>("/api/trees/validate", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify(input)
  });
}

export interface ModelStreamRequest {
  prompt: string;
  system_prompt?: string | null;
  model_id?: string;
  adapter_path?: string | null;
  precision?: "bf16" | "fp16" | "fp32" | "4bit" | "8bit";
  max_new_tokens?: number | null;
  temperature?: number;
  top_p?: number;
  top_k?: number;
  repetition_penalty?: number;
  do_sample?: boolean;
  seed?: number | null;
}

export interface ModelStreamEvent {
  event: string;
  data: Record<string, unknown>;
}

export interface ModelStreamCheckpoint {
  label: string;
  path: string;
  kind: "final" | "checkpoint" | string;
  step?: number | null;
}

export interface ModelStreamModelChoice {
  key: string;
  id: string;
  label: string;
  modelId: string;
  wrapperDir: string;
  settingsPath?: string;
  systemPrompt?: string;
  finalDir: string;
  checkpointDir: string;
  defaultCheckpointPath: string | null;
  checkpoints: ModelStreamCheckpoint[];
}

export interface ModelStreamModelsResponse {
  defaultModelId: string;
  defaultAdapterPath: string | null;
  finalDir: string;
  checkpointDir: string;
  models: ModelStreamModelChoice[];
}

function parseSseMessage(message: string): ModelStreamEvent | null {
  const lines = message.split(/\r?\n/);
  let event = "message";
  const dataLines: string[] = [];
  for (const line of lines) {
    if (line.startsWith("event:")) {
      event = line.slice("event:".length).trim();
    } else if (line.startsWith("data:")) {
      dataLines.push(line.slice("data:".length).trimStart());
    }
  }
  if (dataLines.length === 0) return null;
  const dataText = dataLines.join("\n");
  try {
    return { event, data: JSON.parse(dataText) as Record<string, unknown> };
  } catch {
    return { event, data: { text: dataText } };
  }
}

export async function streamModelGeneration(
  input: ModelStreamRequest,
  onEvent: (event: ModelStreamEvent) => void,
  signal?: AbortSignal
) {
  const response = await fetch("/api/model/stream", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify(input),
    signal
  });
  if (!response.ok || !response.body) {
    const body = await response.json().catch(() => ({ error: response.statusText }));
    throw new Error(String(body.error || response.statusText));
  }

  const reader = response.body.getReader();
  const decoder = new TextDecoder();
  let buffer = "";

  while (true) {
    const { value, done } = await reader.read();
    if (done) break;
    buffer += decoder.decode(value, { stream: true });
    const messages = buffer.split(/\n\n/);
    buffer = messages.pop() ?? "";
    for (const message of messages) {
      const parsed = parseSseMessage(message.trim());
      if (parsed) onEvent(parsed);
    }
  }

  buffer += decoder.decode();
  const parsed = parseSseMessage(buffer.trim());
  if (parsed) onEvent(parsed);
}

export function fetchModelStreamModels() {
  return requestJson<ModelStreamModelsResponse>("/api/model/models");
}