parser.ts

import { readFileSync, writeFileSync } from "fs";

interface Message {
  role: string;
  content: string;
  tool_calls?: string;
  reasoning?: string;
}

interface CurrentFormat {
  messages: Message[];
  meta?: string;
  available_tools?: string;
}

interface Part {
  text?: string;
  fileData?: {
    mimeType: string;
    fileUri: string;
  };
  functionCall?: {
    name: string;
    args: Record<string, any>;
  };
  functionResponse?: {
    name: string;
    response: Record<string, any>;
  };
}

interface Content {
  role: string;
  parts: Part[];
}

interface FunctionDeclaration {
  name: string;
  description?: string;
  parameters?: any; // OpenAPI Schema
  response?: any; // OpenAPI Schema
}

interface Tool {
  functionDeclarations: FunctionDeclaration[];
}

interface GeminiFormat {
  system_instruction?: {
    parts: Part[];
  };
  contents: Content[];
  tools?: Tool[];
}

// Helper function to convert function names to valid format (snake_case, no spaces)
function normalizeToolName(name: string): string {
  return name
    .toLowerCase()
    .replace(/[^a-z0-9_.-]/g, "_")
    .replace(/_+/g, "_")
    .replace(/^[^a-z_]+/, "")
    .substring(0, 64);
}

function normalizeSchemaType(type: any): string {
  // Normalize to standard types then uppercase for Vertex AI
  if (!type || typeof type !== "string") {
    return "STRING"; // Default fallback
  }

  let normalized = type.toLowerCase();
  if (normalized === "dict") normalized = "object";
  else if (normalized === "int" || normalized === "float")
    normalized = "number";
  else if (normalized === "str") normalized = "string";
  else if (normalized === "bool") normalized = "boolean";

  return normalized.toUpperCase();
}

function convertToolsToFunctionDeclarations(
  availableTools: string
): FunctionDeclaration[] {
  try {
    const tools = JSON.parse(availableTools);
    return tools.map((tool: any) => {
      const declaration: FunctionDeclaration = {
        name: normalizeToolName(tool.name),
      };

      if (tool.description) {
        declaration.description = tool.description;
      }

      if (tool.parameters) {
        // Convert the parameters to Vertex AI Schema format
        const params = tool.parameters;

        // Build the schema
        const schema: any = {
          type: normalizeSchemaType(params.type || "object"),
        };

        if (params.properties) {
          // Convert each property to proper Vertex format
          const convertedProperties: any = {};
          for (const [key, prop] of Object.entries(params.properties)) {
            const propValue: any = prop;
            const convertedProp: any = {};

            // Map type field to UPPERCASE Vertex format
            if (propValue.type) {
              convertedProp.type = normalizeSchemaType(propValue.type);
            }

            if (propValue.description) {
              convertedProp.description = propValue.description;
            }

            // NOTE: Remove 'default' field - not supported in Vertex AI Schema
            // if (propValue.default !== undefined) {
            //   convertedProp.default = propValue.default;
            // }

            if (propValue.properties) {
              convertedProp.properties = propValue.properties;
            }

            convertedProperties[key] = convertedProp;
          }
          schema.properties = convertedProperties;
        }

        if (params.required && Array.isArray(params.required)) {
          schema.required = params.required;
        }

        declaration.parameters = schema;
      }

      return declaration;
    });
  } catch (error) {
    console.error("Error parsing available_tools:", error);
    return [];
  }
}

function convertToGeminiFormat(input: CurrentFormat): GeminiFormat {
  const result: GeminiFormat = {
    contents: [],
  };

  // Add tools/function declarations if available
  if (input.available_tools) {
    const functionDeclarations = convertToolsToFunctionDeclarations(
      input.available_tools
    );
    if (functionDeclarations.length > 0) {
      result.tools = [
        {
          functionDeclarations,
        },
      ];
    }
  }

  for (const message of input.messages) {
    // Handle system message
    if (message.role === "system") {
      result.system_instruction = {
        parts: [
          {
            text: message.content,
          },
        ],
      };
      continue;
    }

    // Convert role names to Gemini format
    let geminiRole = message.role;
    if (message.role === "assistant") {
      geminiRole = "model";
    } else if (message.role === "tool") {
      // Tool responses should be "user" role in Gemini
      geminiRole = "user";
    }

    // Build parts array
    const parts: Part[] = [];

    // Handle tool/function responses differently
    if (message.role === "tool" && message.content) {
      // Tool responses need a function name - try to extract from previous context
      // The content could be an ACK or actual response data

      let functionName = "unknown";
      let responseContent: any = {};

      // Try to parse as structured JSON response
      try {
        const responseData = JSON.parse(message.content);
        if (responseData.id) {
          functionName = responseData.id;
          responseContent = responseData;
        } else {
          // Plain JSON but no id, use as response
          responseContent = responseData;
        }
      } catch {
        // Not JSON, treat as plain text in response
        // Need to find the function name from previous tool call
        // Look back in the messages to find the last tool call
        for (let i = input.messages.indexOf(message) - 1; i >= 0; i--) {
          const prevMsg = input.messages[i];
          if (prevMsg.tool_calls) {
            try {
              const toolCalls = JSON.parse(prevMsg.tool_calls);
              if (toolCalls.length > 0) {
                functionName = toolCalls[0].name;
                break;
              }
            } catch {}
          }
        }

        // Structure the response properly
        // Check if it's an ACK or actual data
        if (message.content.includes("<tool_ack")) {
          // ACK - create minimal response
          responseContent = { status: "acknowledged" };
        } else {
          // Regular text response
          responseContent = { result: message.content };
        }
      }

      // Normalize the function name to match declaration format
      const normalizedName = normalizeToolName(functionName);

      parts.push({
        functionResponse: {
          name: normalizedName,
          response: responseContent,
        },
      });
    } else if (message.content) {
      // Add text content for non-tool messages
      // NOTE: Remove reasoning text as per Vertex best practices
      parts.push({
        text: message.content,
      });
    }

    // Parse and add tool_calls as separate functionCall parts
    if (message.tool_calls) {
      try {
        const toolCalls = JSON.parse(message.tool_calls);
        for (const toolCall of toolCalls) {
          parts.push({
            functionCall: {
              name: normalizeToolName(toolCall.name),
              args: toolCall.arguments || {},
            },
          });
        }
      } catch (error) {
        console.error("Error parsing tool_calls:", error);
        // Fallback: add as text if parsing fails
        parts.push({
          text: `Tool Calls: ${message.tool_calls}`,
        });
      }
    }

    // For functionResponse parts, don't include role (Vertex format)
    if (message.role === "tool") {
      result.contents.push({
        parts: parts,
      } as any);
    } else {
      result.contents.push({
        role: geminiRole,
        parts: parts,
      });
    }
  }

  return result;
}

function main() {
  const inputFile = "./train-qwen.jsonl";
  const outputFile = "./train-gemini.jsonl";

  console.log(`Reading from ${inputFile}...`);

  const fileContent = readFileSync(inputFile, "utf-8");
  const lines = fileContent.trim().split("\n");

  console.log(`Found ${lines.length} lines to process`);

  const outputLines: string[] = [];

  for (let i = 0; i < lines.length; i++) {
    try {
      const line = lines[i].trim();
      if (!line) continue;

      const input: CurrentFormat = JSON.parse(line);
      const output = convertToGeminiFormat(input);
      outputLines.push(JSON.stringify(output));

      if ((i + 1) % 100 === 0) {
        console.log(`Processed ${i + 1} lines...`);
      }
    } catch (error) {
      console.error(`Error processing line ${i + 1}:`, error);
    }
  }

  console.log(
    `Writing ${outputLines.length} converted lines to ${outputFile}...`
  );
  writeFileSync(outputFile, outputLines.join("\n") + "\n", "utf-8");

  console.log("✓ Conversion complete!");
  console.log(`Output saved to: ${outputFile}`);
}

main();