moreminimore-vibe/testing/fake-llm-server/chatCompletionHandler.ts

import { Request, Response } from "express";
import fs from "fs";
import path from "path";
import { CANNED_MESSAGE, createStreamChunk } from ".";

let globalCounter = 0;

export const createChatCompletionHandler =
  (prefix: string) => (req: Request, res: Response) => {
    const { stream = false, messages = [] } = req.body;
    console.log("* Received messages", messages);

    // Check if the last message contains "[429]" to simulate rate limiting
    const lastMessage = messages[messages.length - 1];
    if (lastMessage && lastMessage.content === "[429]") {
      return res.status(429).json({
        error: {
          message: "Too many requests. Please try again later.",
          type: "rate_limit_error",
          param: null,
          code: "rate_limit_exceeded",
        },
      });
    }

    let messageContent = CANNED_MESSAGE;
    console.error("LASTMESSAGE", lastMessage);
    // Check if the last message is "[dump]" to write messages to file and return path
    if (
      lastMessage &&
      (Array.isArray(lastMessage.content)
        ? lastMessage.content.some(
            (part: { type: string; text: string }) =>
              part.type === "text" && part.text.includes("[dump]"),
          )
        : lastMessage.content.includes("[dump]"))
    ) {
      const timestamp = Date.now();
      const generatedDir = path.join(__dirname, "generated");

      // Create generated directory if it doesn't exist
      if (!fs.existsSync(generatedDir)) {
        fs.mkdirSync(generatedDir, { recursive: true });
      }

      const dumpFilePath = path.join(generatedDir, `${timestamp}.json`);

      try {
        fs.writeFileSync(
          dumpFilePath,
          JSON.stringify(
            {
              body: req.body,
              headers: { authorization: req.headers["authorization"] },
            },
            null,
            2,
          ).replace(/\r\n/g, "\n"),
          "utf-8",
        );
        console.log(`* Dumped messages to: ${dumpFilePath}`);
        messageContent = `[[dyad-dump-path=${dumpFilePath}]]`;
      } catch (error) {
        console.error(`* Error writing dump file: ${error}`);
        messageContent = `Error: Could not write dump file: ${error}`;
      }
    }

    if (lastMessage && lastMessage.content === "[increment]") {
      globalCounter++;
      messageContent = `counter=${globalCounter}`;
    }

    // Check if the last message starts with "tc=" to load test case file
    if (
      lastMessage &&
      lastMessage.content &&
      typeof lastMessage.content === "string" &&
      lastMessage.content.startsWith("tc=")
    ) {
      const testCaseName = lastMessage.content.slice(3); // Remove "tc=" prefix
      const testFilePath = path.join(
        __dirname,
        "..",
        "..",
        "..",
        "e2e-tests",
        "fixtures",
        prefix,
        `${testCaseName}.md`,
      );

      try {
        if (fs.existsSync(testFilePath)) {
          messageContent = fs.readFileSync(testFilePath, "utf-8");
          console.log(`* Loaded test case: ${testCaseName}`);
        } else {
          console.log(`* Test case file not found: ${testFilePath}`);
          messageContent = `Error: Test case file not found: ${testCaseName}.md`;
        }
      } catch (error) {
        console.error(`* Error reading test case file: ${error}`);
        messageContent = `Error: Could not read test case file: ${testCaseName}.md`;
      }
    }

    // Non-streaming response
    if (!stream) {
      return res.json({
        id: `chatcmpl-${Date.now()}`,
        object: "chat.completion",
        created: Math.floor(Date.now() / 1000),
        model: "fake-model",
        choices: [
          {
            index: 0,
            message: {
              role: "assistant",
              content: messageContent,
            },
            finish_reason: "stop",
          },
        ],
      });
    }

    // Streaming response
    res.setHeader("Content-Type", "text/event-stream");
    res.setHeader("Cache-Control", "no-cache");
    res.setHeader("Connection", "keep-alive");

    // Split the message into characters to simulate streaming
    const message = messageContent;
    const messageChars = message.split("");

    // Stream each character with a delay
    let index = 0;
    const batchSize = 8;

    // Send role first
    res.write(createStreamChunk("", "assistant"));

    const interval = setInterval(() => {
      if (index < messageChars.length) {
        // Get the next batch of characters (up to batchSize)
        const batch = messageChars.slice(index, index + batchSize).join("");
        res.write(createStreamChunk(batch));
        index += batchSize;
      } else {
        // Send the final chunk
        res.write(createStreamChunk("", "assistant", true));
        clearInterval(interval);
        res.end();
      }
    }, 10);
  };