Neon-Desk/lib/server/research-copilot-format.ts

import { randomUUID } from 'node:crypto';
import type {
  ResearchCopilotSuggestedAction,
  ResearchMemoSection,
  SearchResult
} from '@/lib/types';

type ParsedCopilotPayload = {
  answerMarkdown: string;
  followUps: string[];
  suggestedActions: ResearchCopilotSuggestedAction[];
  citationIndexes: number[];
};

const MAX_FOLLOW_UPS = 4;
const MAX_SUGGESTED_ACTIONS = 3;

function truncate(value: string, maxLength: number) {
  const normalized = value.trim();
  if (normalized.length <= maxLength) {
    return normalized;
  }

  return `${normalized.slice(0, maxLength - 1).trimEnd()}…`;
}

function buildSessionTitle(query: string) {
  return truncate(query, 72);
}

export function extractJsonObject(text: string) {
  const fenced = text.match(/```json\s*([\s\S]*?)```/i)?.[1];
  if (fenced) {
    return fenced.trim();
  }

  const start = text.indexOf('{');
  const end = text.lastIndexOf('}');
  if (start >= 0 && end > start) {
    return text.slice(start, end + 1).trim();
  }

  return null;
}

export function parseCitationIndexes(value: string, evidenceLength: number) {
  const matches = [...value.matchAll(/\[(\d+)\]/g)];
  const seen = new Set<number>();
  const indexes: number[] = [];

  for (const match of matches) {
    const parsed = Number(match[1]);
    if (!Number.isInteger(parsed) || parsed < 1 || parsed > evidenceLength || seen.has(parsed)) {
      continue;
    }

    seen.add(parsed);
    indexes.push(parsed);
  }

  return indexes;
}

function parseStringArray(value: unknown, maxItems: number) {
  if (!Array.isArray(value)) {
    return [];
  }

  return value
    .filter((entry): entry is string => typeof entry === 'string' && entry.trim().length > 0)
    .map((entry) => truncate(entry, 220))
    .slice(0, maxItems);
}

function normalizeSuggestedAction(
  value: unknown,
  fallbackQuery: string
): ResearchCopilotSuggestedAction | null {
  if (!value || typeof value !== 'object' || Array.isArray(value)) {
    return null;
  }

  const candidate = value as Record<string, unknown>;
  const type = candidate.type;
  if (type !== 'draft_note' && type !== 'draft_memo_section' && type !== 'queue_research_brief') {
    return null;
  }

  const label = typeof candidate.label === 'string' && candidate.label.trim().length > 0
    ? truncate(candidate.label, 80)
    : type === 'draft_note'
      ? 'Use as note draft'
      : type === 'draft_memo_section'
        ? 'Use as memo draft'
        : 'Queue research brief';

  const section = candidate.section === 'thesis'
    || candidate.section === 'variant_view'
    || candidate.section === 'catalysts'
    || candidate.section === 'risks'
    || candidate.section === 'disconfirming_evidence'
    || candidate.section === 'next_actions'
    ? candidate.section
    : null;
  const description = typeof candidate.description === 'string' && candidate.description.trim().length > 0
    ? truncate(candidate.description, 180)
    : null;
  const title = typeof candidate.title === 'string' && candidate.title.trim().length > 0
    ? truncate(candidate.title, 120)
    : null;
  const contentMarkdown = typeof candidate.contentMarkdown === 'string' && candidate.contentMarkdown.trim().length > 0
    ? candidate.contentMarkdown.trim()
    : null;
  const citationIndexes = Array.isArray(candidate.citationIndexes)
    ? candidate.citationIndexes
      .map((entry) => Math.trunc(Number(entry)))
      .filter((entry, index, source) => Number.isInteger(entry) && entry > 0 && source.indexOf(entry) === index)
    : [];
  const query = typeof candidate.query === 'string' && candidate.query.trim().length > 0
    ? truncate(candidate.query, 180)
    : type === 'queue_research_brief'
      ? fallbackQuery
      : null;

  if ((type === 'draft_note' || type === 'draft_memo_section') && !contentMarkdown) {
    return null;
  }

  if (type === 'draft_memo_section' && !section) {
    return null;
  }

  return {
    id: randomUUID(),
    type,
    label,
    description,
    section,
    title,
    content_markdown: contentMarkdown,
    citation_indexes: citationIndexes,
    query
  };
}

function buildFallbackActions(query: string, memoSection: ResearchMemoSection | null, answerMarkdown: string) {
  return [
    {
      id: randomUUID(),
      type: memoSection ? 'draft_memo_section' : 'draft_note',
      label: memoSection ? 'Use as memo draft' : 'Use as note draft',
      description: memoSection
        ? `Populate ${memoSection.replace('_', ' ')} with this answer for review.`
        : 'Populate the note draft editor with this answer for review.',
      section: memoSection,
      title: memoSection ? null : buildSessionTitle(query),
      content_markdown: answerMarkdown,
      citation_indexes: [],
      query: null
    },
    {
      id: randomUUID(),
      type: 'queue_research_brief',
      label: 'Queue research brief',
      description: 'Run a background synthesis job and save a longer-form brief to the library.',
      section: null,
      title: null,
      content_markdown: null,
      citation_indexes: [],
      query
    }
  ] satisfies ResearchCopilotSuggestedAction[];
}

export function parseCopilotResponse(
  rawText: string,
  evidence: SearchResult[],
  query: string,
  memoSection: ResearchMemoSection | null
): ParsedCopilotPayload {
  const jsonText = extractJsonObject(rawText);
  if (!jsonText) {
    const answerMarkdown = rawText.trim() || 'Insufficient evidence to answer from the indexed sources.';
    return {
      answerMarkdown,
      followUps: [],
      suggestedActions: buildFallbackActions(query, memoSection, answerMarkdown),
      citationIndexes: evidence.slice(0, Math.min(3, evidence.length)).map((_value, index) => index + 1)
    };
  }

  try {
    const parsed = JSON.parse(jsonText) as Record<string, unknown>;
    const answerMarkdown = typeof parsed.answerMarkdown === 'string' && parsed.answerMarkdown.trim().length > 0
      ? parsed.answerMarkdown.trim()
      : 'Insufficient evidence to answer from the indexed sources.';
    const citationIndexes = parseCitationIndexes(answerMarkdown, evidence.length);
    const followUps = parseStringArray(parsed.followUps, MAX_FOLLOW_UPS);
    const suggestedActions = Array.isArray(parsed.suggestedActions)
      ? parsed.suggestedActions
        .map((entry) => normalizeSuggestedAction(entry, query))
        .filter((entry): entry is ResearchCopilotSuggestedAction => Boolean(entry))
        .slice(0, MAX_SUGGESTED_ACTIONS)
      : [];

    return {
      answerMarkdown,
      followUps,
      suggestedActions: suggestedActions.length > 0
        ? suggestedActions
        : buildFallbackActions(query, memoSection, answerMarkdown),
      citationIndexes: citationIndexes.length > 0
        ? citationIndexes
        : evidence.slice(0, Math.min(3, evidence.length)).map((_value, index) => index + 1)
    };
  } catch {
    const answerMarkdown = rawText.trim() || 'Insufficient evidence to answer from the indexed sources.';
    return {
      answerMarkdown,
      followUps: [],
      suggestedActions: buildFallbackActions(query, memoSection, answerMarkdown),
      citationIndexes: evidence.slice(0, Math.min(3, evidence.length)).map((_value, index) => index + 1)
    };
  }
}