socaichat/langgraph/src/agent/writer/nodes/tool-executor.ts

/**
 * Writer tool executor node: executes doc_create / doc_edit / doc_translate,
 * pushes Gen-UI canvas-doc cards, and returns ToolMessages.
 */
import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
import type ComponentMap from "../../../agent-uis/index.js";
import { LangGraphRunnableConfig } from "@langchain/langgraph";
import { AIMessage } from "@langchain/core/messages";
import { createLlm, type ModelMode } from "@/agent/utils/create-llm";
import { WriterState, WriterUpdate } from "../types.js";
import {
  docCreateSchema,
  docEditSchema,
  docTranslateSchema,
} from "./tool-defs.js";
import { formatToolError } from "@/agent/utils/retry";

/**
 * Search backwards through messages to find the most recent canvas-doc content
 * for a given doc_id. Returns the content string or undefined.
 */
function findDocContent(
  state: WriterState,
  docId: string,
): { title: string; content: string; type: "markdown" | "code"; language?: string } | undefined {
  // Look through tool messages for a doc_create result containing this doc_id
  for (let i = state.messages.length - 1; i >= 0; i--) {
    const msg = state.messages[i];
    if (msg._getType() === "tool") {
      try {
        const parsed = JSON.parse(
          typeof msg.content === "string" ? msg.content : "",
        );
        if (parsed.doc_id === docId && parsed._doc_content) {
          return {
            title: parsed.title ?? "文档",
            content: parsed._doc_content,
            type: parsed.type ?? "markdown",
            language: parsed.language,
          };
        }
      } catch {
        // not JSON, skip
      }
    }
  }
  return undefined;
}

export async function writerToolExecutorNode(
  state: WriterState,
  config: LangGraphRunnableConfig,
): Promise<WriterUpdate> {
  const ui = typedUi<typeof ComponentMap>(config);
  const modelMode = ((config.configurable?.modelMode as string) ?? "auto") as ModelMode;

  const lastAiMessage = [...state.messages]
    .reverse()
    .find(
      (m): m is AIMessage =>
        (m as AIMessage).tool_calls !== undefined &&
        ((m as AIMessage).tool_calls?.length ?? 0) > 0,
    ) as AIMessage | undefined;

  if (!lastAiMessage?.tool_calls?.length) {
    return { ui: ui.items, timestamp: Date.now() };
  }

  const toolCalls = lastAiMessage.tool_calls!;
  const toolMessages: Array<{
    role: "tool";
    tool_call_id: string;
    content: string;
  }> = [];

  for (const tc of toolCalls) {
    const name = tc.name;
    const args = tc.args;
    const id = tc.id ?? "";

    try {
      switch (name) {
        case "doc_create": {
          const parsed = docCreateSchema.parse(args);
          const docId = `doc-${Date.now()}`;

          ui.push(
            {
              name: "canvas-doc",
              props: {
                doc_id: docId,
                title: parsed.title,
                content: parsed.content,
                type: parsed.type,
                language: parsed.language ?? "",
              },
            },
            { message: lastAiMessage },
          );

          toolMessages.push({
            role: "tool",
            tool_call_id: id,
            content: JSON.stringify({
              doc_id: docId,
              title: parsed.title,
              type: parsed.type,
              language: parsed.language,
              char_count: parsed.content.length,
              // Store content for later edits (not shown to user in text)
              _doc_content: parsed.content,
            }),
          });
          break;
        }

        case "doc_edit": {
          const parsed = docEditSchema.parse(args);
          const existing = findDocContent(state, parsed.doc_id);

          const editLlm = createLlm({ modelMode, maxTokens: 8192 });
          const originalContent = existing?.content ?? "(原文档内容不可用)";
          const originalTitle = existing?.title ?? "文档";
          const docType = existing?.type ?? "markdown";

          const editResult = await editLlm.invoke([
            {
              role: "system",
              content:
                "你是文档编辑助手。根据用户指令修改以下文档，直接输出修改后的完整文档内容，不要添加任何解释或说明。",
            },
            {
              role: "user",
              content: `原文档标题: ${originalTitle}\n原文档内容:\n${originalContent}\n\n编辑指令: ${parsed.instructions}\n\n请直接输出修改后的完整文档内容：`,
            },
          ]);

          const editedContent =
            typeof editResult.content === "string"
              ? editResult.content
              : "";

          ui.push(
            {
              name: "canvas-doc",
              props: {
                doc_id: parsed.doc_id,
                title: originalTitle,
                content: editedContent,
                type: docType,
                language: existing?.language ?? "",
              },
            },
            { message: lastAiMessage },
          );

          toolMessages.push({
            role: "tool",
            tool_call_id: id,
            content: JSON.stringify({
              doc_id: parsed.doc_id,
              title: originalTitle,
              type: docType,
              char_count: editedContent.length,
              _doc_content: editedContent,
            }),
          });
          break;
        }

        case "doc_translate": {
          const parsed = docTranslateSchema.parse(args);
          const existing = findDocContent(state, parsed.doc_id);

          const translateLlm = createLlm({ modelMode, maxTokens: 8192 });
          const originalContent = existing?.content ?? "(原文档内容不可用)";
          const originalTitle = existing?.title ?? "文档";
          const docType = existing?.type ?? "markdown";

          const translateResult = await translateLlm.invoke([
            {
              role: "system",
              content: `你是专业翻译。将以下文档翻译为${parsed.target_language}，保持原始格式和结构，直接输出翻译后的完整文档内容。`,
            },
            {
              role: "user",
              content: originalContent,
            },
          ]);

          const translatedContent =
            typeof translateResult.content === "string"
              ? translateResult.content
              : "";

          const translatedTitle = `${originalTitle} (${parsed.target_language})`;

          ui.push(
            {
              name: "canvas-doc",
              props: {
                doc_id: parsed.doc_id,
                title: translatedTitle,
                content: translatedContent,
                type: docType,
                language: existing?.language ?? "",
              },
            },
            { message: lastAiMessage },
          );

          toolMessages.push({
            role: "tool",
            tool_call_id: id,
            content: JSON.stringify({
              doc_id: parsed.doc_id,
              title: translatedTitle,
              type: docType,
              char_count: translatedContent.length,
              _doc_content: translatedContent,
            }),
          });
          break;
        }

        default:
          toolMessages.push({
            role: "tool",
            tool_call_id: id,
            content: `Unknown tool: ${name}`,
          });
      }
    } catch (e) {
      toolMessages.push({
        role: "tool",
        tool_call_id: id,
        content: formatToolError(name, e),
      });
    }
  }

  return {
    messages: toolMessages,
    ui: ui.items,
    timestamp: Date.now(),
  };
}