TJWaterFrontend_Refine/src/components/chat/chatMessageSections.ts

export type AssistantMessageSections = {
  answer: string;
  thought: string | null;
  thoughtComplete: boolean;
};

/* ------------------------------------------------------------------ */
/*  Tool-call types                                                   */
/* ------------------------------------------------------------------ */

export type ToolCall = {
  id: string;
  tool: string;
  params: Record<string, unknown>;
};

export type ContentSegment =
  | { type: "text"; content: string }
  | { type: "tool_call"; toolCall: ToolCall }
  | { type: "tool_call_pending" };

export type ParsedToolContent = {
  segments: ContentSegment[];
  toolCalls: ToolCall[];
};

/* ------------------------------------------------------------------ */
/*  Think-block parsing                                               */
/* ------------------------------------------------------------------ */

const THINK_BLOCK_PATTERN = /<think>([\s\S]*?)<\/think>/gi;
const THINK_OPEN_TAG = "<think>";
const THINK_CLOSE_TAG = "</think>";

export const parseAssistantMessageSections = (
  content: string,
): AssistantMessageSections => {
  if (!content) {
    return { answer: "", thought: null, thoughtComplete: false };
  }

  const thoughtParts: string[] = [];
  let answer = content;

  answer = answer.replace(THINK_BLOCK_PATTERN, (_, thoughtContent: string) => {
    const trimmedThought = thoughtContent.trim();
    if (trimmedThought) {
      thoughtParts.push(trimmedThought);
    }

    return "\n";
  });

  const lastOpenIndex = answer.lastIndexOf(THINK_OPEN_TAG);
  const lastCloseIndex = answer.lastIndexOf(THINK_CLOSE_TAG);
  const hasUnclosedThought =
    lastOpenIndex !== -1 && lastOpenIndex > lastCloseIndex;

  if (hasUnclosedThought) {
    const streamingThought = answer
      .slice(lastOpenIndex + THINK_OPEN_TAG.length)
      .trim();

    if (streamingThought) {
      thoughtParts.push(streamingThought);
    }

    answer = answer.slice(0, lastOpenIndex);
  }

  const normalizedAnswer = answer.replace(/\n{3,}/g, "\n\n").trim();
  const normalizedThought = thoughtParts.join("\n\n").trim();

  return {
    answer: normalizedAnswer,
    thought: normalizedThought || null,
    thoughtComplete: Boolean(normalizedThought) && !hasUnclosedThought,
  };
};

/* ------------------------------------------------------------------ */
/*  Tool-call parsing                                                 */
/*                                                                    */
/*  AI responses may embed tool calls using:                          */
/*    <tool_call>{"tool":"locate_pipes","params":{...}}</tool_call>   */
/*                                                                    */
/*  Returns ordered segments (text + tool_call interleaved) so the    */
/*  UI can render them inline where the AI placed them.               */
/* ------------------------------------------------------------------ */

const TOOL_CALL_BLOCK_PATTERN = /<tool_call>([\s\S]*?)<\/tool_call>/gi;
const TOOL_CALL_OPEN_TAG = "<tool_call>";

/** Regex to strip partial opening tag at the end of text during streaming. */
const PARTIAL_TOOL_TAG_TAIL = /<(?:t(?:o(?:o(?:l(?:_(?:c(?:a(?:l(?:l)?)?)?)?)?)?)?)?)?$/;

export const parseContentWithToolCalls = (
  content: string,
): ParsedToolContent => {
  if (!content) {
    return { segments: [], toolCalls: [] };
  }

  const segments: ContentSegment[] = [];
  const toolCalls: ToolCall[] = [];
  let lastIndex = 0;
  let tcIndex = 0;

  // Find all complete <tool_call>...</tool_call> blocks
  const regex = /<tool_call>([\s\S]*?)<\/tool_call>/gi;
  let match: RegExpExecArray | null;

  while ((match = regex.exec(content)) !== null) {
    // Text before this tool call
    const textBefore = content.slice(lastIndex, match.index);
    if (textBefore.trim()) {
      segments.push({ type: "text", content: textBefore.trim() });
    }

    // Parse the tool call JSON
    try {
      const parsed = JSON.parse(match[1].trim()) as {
        tool?: string;
        params?: Record<string, unknown>;
      };
      const toolCall: ToolCall = {
        id: `tc-${tcIndex++}`,
        tool: parsed.tool ?? "unknown",
        params: parsed.params ?? {},
      };
      segments.push({ type: "tool_call", toolCall });
      toolCalls.push(toolCall);
    } catch {
      // Malformed JSON – treat as plain text
      segments.push({ type: "text", content: match[0] });
    }

    lastIndex = match.index + match[0].length;
  }

  // Handle remaining text after the last match
  const remaining = content.slice(lastIndex);

  // Check for an unclosed <tool_call> tag (still streaming)
  const unclosedIdx = remaining.lastIndexOf(TOOL_CALL_OPEN_TAG);
  if (unclosedIdx !== -1) {
    const textBefore = remaining.slice(0, unclosedIdx);
    if (textBefore.trim()) {
      segments.push({ type: "text", content: textBefore.trim() });
    }
    segments.push({ type: "tool_call_pending" });
  } else {
    // Strip partial opening tags at the end (e.g. "<tool_c" while streaming)
    const cleaned = remaining.replace(PARTIAL_TOOL_TAG_TAIL, "").trim();
    if (cleaned) {
      segments.push({ type: "text", content: cleaned });
    }
  }

  // If nothing was parsed, return the original content as a single text segment
  if (segments.length === 0) {
    segments.push({ type: "text", content });
  }

  return { segments, toolCalls };
};