285 lines · 9.5 KB · TypeScript
RawDownload
// CRITICALimport type { AgentMessage } from "@mariozechner/pi-agent-core";import type {  AssistantMessage,  ImageContent,  Message,  Model,  TextContent,  ThinkingContent,  ToolCall,  ToolResultMessage,  Usage,} from "@mariozechner/pi-ai"; type StoredMessageRecord = Record<string, unknown>; const getString = (value: unknown): string | undefined => (typeof value === "string" ? value : undefined);const getNumber = (value: unknown): number | undefined =>  typeof value === "number" && !Number.isNaN(value) ? value : undefined; const parseJson = (value: unknown): unknown => {  if (typeof value !== "string") return value;  try {    return JSON.parse(value);  } catch {    return value;  }}; const parseTimestamp = (value: unknown): number => {  if (typeof value === "number") return value;  if (typeof value === "string") {    const parsed = Date.parse(value);    if (!Number.isNaN(parsed)) return parsed;  }  return Date.now();}; const normalizeToolArguments = (value: unknown): Record<string, unknown> => {  if (!value) return {};  if (typeof value === "object" && !Array.isArray(value)) {    return value as Record<string, unknown>;  }  if (typeof value === "string") {    try {      const parsed = JSON.parse(value);      if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {        return parsed as Record<string, unknown>;      }    } catch {      return { raw: value };    }  }  return { value };}; const extractToolResult = (  raw: unknown,): { text: string; isError: boolean; details: Record<string, unknown> } | null => {  if (raw === undefined || raw === null) return null;  if (typeof raw === "string") {    return { text: raw, isError: false, details: { content: raw } };  }  if (typeof raw === "object") {    const record = raw as Record<string, unknown>;    const isError = record["isError"] === true;    const content = record["content"] ?? record["result"] ?? record["output"];    if (typeof content === "string") {      return { text: content, isError, details: record };    }    try {      return { text: JSON.stringify(content ?? record), isError, details: record };    } catch {      return { text: String(content ?? raw), isError, details: record };    }  }  return { text: String(raw), isError: false, details: { content: raw } };}; const buildUsage = (message: StoredMessageRecord): Usage => {  const promptTokens = getNumber(message["prompt_tokens"]) ?? getNumber(message["request_prompt_tokens"]) ?? 0;  const completionTokens = getNumber(message["completion_tokens"]) ?? getNumber(message["request_completion_tokens"]) ?? 0;  const totalTokens =    getNumber(message["total_tokens"]) ??    getNumber(message["request_total_input_tokens"]) ??    promptTokens + completionTokens;  return {    input: promptTokens,    output: completionTokens,    cacheRead: 0,    cacheWrite: 0,    totalTokens,    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },  };}; const buildAssistantContent = (  message: StoredMessageRecord,  toolCallsById: Map<string, ToolCall>,): Array<TextContent | ThinkingContent | ToolCall> => {  const contentBlocks: Array<TextContent | ThinkingContent | ToolCall> = [];  const rawParts = Array.isArray(message["parts"]) ? (message["parts"] as Array<Record<string, unknown>>) : [];  const addedToolCalls = new Set<string>();   for (const part of rawParts) {    const type = getString(part["type"]) ?? "";    if (type === "text") {      const text = getString(part["text"]) ?? "";      if (text) {        contentBlocks.push({ type: "text", text });      }      continue;    }    if (type === "reasoning") {      const thinking = getString(part["text"]) ?? "";      if (thinking) {        contentBlocks.push({ type: "thinking", thinking });      }      continue;    }    const isToolPart = type === "dynamic-tool" || type.startsWith("tool-");    if (isToolPart) {      const toolCallId = getString(part["toolCallId"]) ?? "";      const toolName = getString(part["toolName"]) ?? (type.startsWith("tool-") ? type.replace(/^tool-/, "") : "tool");      const input = normalizeToolArguments(part["input"]);      if (toolCallId) {        const existing = toolCallsById.get(toolCallId);        if (existing) {          contentBlocks.push(existing);        } else {          contentBlocks.push({ type: "toolCall", id: toolCallId, name: toolName, arguments: input });        }        addedToolCalls.add(toolCallId);      }      continue;    }  }   for (const [toolCallId, toolCall] of toolCallsById.entries()) {    if (!addedToolCalls.has(toolCallId)) {      contentBlocks.push(toolCall);    }  }   if (contentBlocks.length === 0) {    const content = getString(message["content"]) ?? "";    if (content) {      contentBlocks.push({ type: "text", text: content });    }  }   return contentBlocks;}; const buildToolCalls = (message: StoredMessageRecord): Map<string, ToolCall> => {  const calls = Array.isArray(message["tool_calls"]) ? message["tool_calls"] : parseJson(message["tool_calls"]);  const toolCalls: unknown[] = Array.isArray(calls) ? calls : [];  const map = new Map<string, ToolCall>();   for (const call of toolCalls) {    if (!call || typeof call !== "object") continue;    const record = call as Record<string, unknown>;    const id = getString(record["id"]) ?? "";    const functionPayload = record["function"] as Record<string, unknown> | undefined;    const name = getString(functionPayload?.["name"]) ?? "tool";    const argsRaw = functionPayload?.["arguments"];    const args = normalizeToolArguments(argsRaw);    if (!id) continue;    map.set(id, { type: "toolCall", id, name, arguments: args });  }   return map;}; const buildToolResults = (message: StoredMessageRecord): ToolResultMessage[] => {  const calls = Array.isArray(message["tool_calls"]) ? message["tool_calls"] : parseJson(message["tool_calls"]);  const toolCallsArray: unknown[] = Array.isArray(calls) ? calls : [];  const results: ToolResultMessage[] = [];  const agentFsTools = new Set([    "list_files",    "read_file",    "write_file",    "delete_file",    "make_directory",    "move_file",  ]);   for (const call of toolCallsArray) {    if (!call || typeof call !== "object") continue;    const record = call as Record<string, unknown>;    const id = getString(record["id"]) ?? "";    const functionPayload = record["function"] as Record<string, unknown> | undefined;    const name = getString(functionPayload?.["name"]) ?? "tool";    if (agentFsTools.has(name)) {      results.push({        role: "toolResult",        toolCallId: id,        toolName: name,        content: [{ type: "text", text: "[completed]" }],        details: {},        isError: false,        timestamp: parseTimestamp(record["created_at"] ?? message["created_at"]),      });      continue;    }    const result = extractToolResult(record["result"]);    if (!id || !result) continue;    results.push({      role: "toolResult",      toolCallId: id,      toolName: name,      content: [{ type: "text", text: result.text }],      details: result.details,      isError: result.isError,      timestamp: parseTimestamp(record["created_at"] ?? message["created_at"]),    });  }   return results;}; export const mapStoredMessagesToAgentMessages = (  storedMessages: StoredMessageRecord[],  fallbackModel: Model<"openai-completions">,): AgentMessage[] => {  const mapped: AgentMessage[] = [];   for (const message of storedMessages) {    const role = getString(message["role"]) ?? "assistant";    const timestamp = parseTimestamp(message["created_at"]);     if (role === "user") {      const rawParts = Array.isArray(message["parts"]) ? (message["parts"] as Array<Record<string, unknown>>) : [];      const blocks: Array<TextContent | ImageContent> = [];       for (const part of rawParts) {        const type = getString(part["type"]) ?? "";        if (type === "text") {          const text = getString(part["text"]) ?? "";          if (text) blocks.push({ type: "text", text });          continue;        }        if (type === "image") {          const data = getString(part["data"]) ?? "";          const mimeType = getString(part["mimeType"]) ?? "";          if (data && mimeType) blocks.push({ type: "image", data, mimeType });        }      }       const hasImages = blocks.some((block) => block.type === "image");      const fallback = getString(message["content"]) ?? "";      const content = hasImages ? blocks : blocks.length > 0 ? blocks.map((b) => (b.type === "text" ? b.text : "")).join("") : fallback;      mapped.push({ role: "user", content, timestamp });      continue;    }     if (role === "assistant") {      const toolCalls = buildToolCalls(message);      const content = buildAssistantContent(message, toolCalls);      const usage = buildUsage(message);      const stopReason = content.some((block) => block.type === "toolCall") ? "toolUse" : "stop";      const assistant: AssistantMessage = {        role: "assistant",        content,        api: fallbackModel.api,        provider: fallbackModel.provider,        model: getString(message["model"]) ?? fallbackModel.id,        usage,        stopReason,        timestamp,      };      mapped.push(assistant);      const toolResults = buildToolResults(message);      if (toolResults.length > 0) {        mapped.push(...toolResults);      }      continue;    }  }   return mapped;}; export const mapAgentMessagesToLlm = (messages: AgentMessage[]): Message[] =>  messages.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[];