copilotkit--copilotkit/packages/runtime/tests/service-adapters/openai/openai-adapter.test.ts

// Mock the modules first
vi.mock("openai", () => {
  function MockOpenAI() {}
  return { default: MockOpenAI };
});

// Mock the OpenAIAdapter class to avoid the "new OpenAI()" issue
vi.mock("../../../src/service-adapters/openai/openai-adapter", () => {
  class MockOpenAIAdapter {
    _openai: any;
    model: string = "gpt-4o";
    keepSystemRole: boolean = false;
    disableParallelToolCalls: boolean = false;

    constructor() {
      this._openai = {
        beta: {
          chat: {
            completions: {
              stream: vi.fn(),
            },
          },
        },
      };
    }

    get openai() {
      return this._openai;
    }

    async process(request: any) {
      // Mock implementation that calls our event source but doesn't do the actual processing
      request.eventSource.stream(async (stream: any) => {
        stream.complete();
        return Promise.resolve();
      });

      return { threadId: request.threadId || "mock-thread-id" };
    }
  }

  return { OpenAIAdapter: MockOpenAIAdapter };
});

// Mock the Message classes since they use TypeGraphQL decorators
vi.mock("../../../src/graphql/types/converted", () => {
  // Create minimal implementations of the message classes
  class MockTextMessage {
    content: string;
    role: string;
    id: string;

    constructor(options: { role: string; content: string }) {
      this.role = options.role;
      this.content = options.content;
      this.id = "mock-text-" + Math.random().toString(36).slice(7);
    }

    isTextMessage() {
      return true;
    }
    isImageMessage() {
      return false;
    }
    isActionExecutionMessage() {
      return false;
    }
    isResultMessage() {
      return false;
    }
  }

  class MockActionExecutionMessage {
    id: string;
    name: string;
    arguments: string;

    constructor(params: { id: string; name: string; arguments: string }) {
      this.id = params.id;
      this.name = params.name;
      this.arguments = params.arguments;
    }

    isTextMessage() {
      return false;
    }
    isImageMessage() {
      return false;
    }
    isActionExecutionMessage() {
      return true;
    }
    isResultMessage() {
      return false;
    }
  }

  class MockResultMessage {
    actionExecutionId: string;
    result: string;
    id: string;

    constructor(params: { actionExecutionId: string; result: string }) {
      this.actionExecutionId = params.actionExecutionId;
      this.result = params.result;
      this.id = "mock-result-" + Math.random().toString(36).slice(7);
    }

    isTextMessage() {
      return false;
    }
    isImageMessage() {
      return false;
    }
    isActionExecutionMessage() {
      return false;
    }
    isResultMessage() {
      return true;
    }
  }

  return {
    TextMessage: MockTextMessage,
    ActionExecutionMessage: MockActionExecutionMessage,
    ResultMessage: MockResultMessage,
    Role: {
      assistant: "assistant",
      developer: "developer",
      system: "system",
      tool: "tool",
      user: "user",
    },
  };
});

// Now import the modules (vi.mock is hoisted above imports, so these get the mocked versions)
import { OpenAIAdapter } from "../../../src/service-adapters/openai/openai-adapter";
import {
  TextMessage,
  ActionExecutionMessage,
  ResultMessage,
  Role,
} from "../../../src/graphql/types/converted";

describe("OpenAIAdapter", () => {
  let adapter: OpenAIAdapter;
  let mockEventSource: any;

  beforeEach(() => {
    vi.clearAllMocks();
    adapter = new OpenAIAdapter();
    mockEventSource = {
      stream: vi.fn((callback) => {
        const mockStream = {
          sendTextMessageStart: vi.fn(),
          sendTextMessageContent: vi.fn(),
          sendTextMessageEnd: vi.fn(),
          sendActionExecutionStart: vi.fn(),
          sendActionExecutionArgs: vi.fn(),
          sendActionExecutionEnd: vi.fn(),
          complete: vi.fn(),
        };
        callback(mockStream);
      }),
    };
  });

  describe("Tool ID handling", () => {
    it("should filter out tool_result messages that don't have corresponding tool_call IDs", async () => {
      // Create messages including one valid pair and one invalid tool_result
      const systemMessage = new TextMessage({
        role: Role.system,
        content: "System message",
      });
      const userMessage = new TextMessage({
        role: Role.user,
        content: "User message",
      });

      // Valid tool execution message
      const validToolExecution = new ActionExecutionMessage({
        id: "valid-tool-id",
        name: "validTool",
        arguments: '{"arg":"value"}',
      });

      // Valid result for the above tool
      const validToolResult = new ResultMessage({
        actionExecutionId: "valid-tool-id",
        result: '{"result":"success"}',
      });

      // Invalid tool result with no corresponding tool execution
      const invalidToolResult = new ResultMessage({
        actionExecutionId: "invalid-tool-id",
        result: '{"result":"failure"}',
      });

      // Spy on the process method to test it's called properly
      const processSpy = vi.spyOn(adapter, "process");

      await adapter.process({
        threadId: "test-thread",
        model: "gpt-4o",
        messages: [
          systemMessage,
          userMessage,
          validToolExecution,
          validToolResult,
          invalidToolResult,
        ],
        actions: [],
        eventSource: mockEventSource,
        forwardedParameters: {},
      });

      // Verify the process method was called
      expect(processSpy).toHaveBeenCalledTimes(1);

      // Verify the stream function was called
      expect(mockEventSource.stream).toHaveBeenCalled();
    });

    it("should handle duplicate tool IDs by only using each once", async () => {
      // Create messages including duplicate tool results for the same ID
      const systemMessage = new TextMessage({
        role: Role.system,
        content: "System message",
      });

      // Valid tool execution message
      const toolExecution = new ActionExecutionMessage({
        id: "tool-id-1",
        name: "someTool",
        arguments: '{"arg":"value"}',
      });

      // Two results for the same tool ID
      const firstToolResult = new ResultMessage({
        actionExecutionId: "tool-id-1",
        result: '{"result":"first"}',
      });

      const duplicateToolResult = new ResultMessage({
        actionExecutionId: "tool-id-1",
        result: '{"result":"duplicate"}',
      });

      // Spy on the process method to test it's called properly
      const processSpy = vi.spyOn(adapter, "process");

      await adapter.process({
        threadId: "test-thread",
        model: "gpt-4o",
        messages: [
          systemMessage,
          toolExecution,
          firstToolResult,
          duplicateToolResult,
        ],
        actions: [],
        eventSource: mockEventSource,
        forwardedParameters: {},
      });

      // Verify the process method was called
      expect(processSpy).toHaveBeenCalledTimes(1);

      // Verify the stream function was called
      expect(mockEventSource.stream).toHaveBeenCalled();
    });

    it("should call the stream method on eventSource", async () => {
      // Create messages
      const systemMessage = new TextMessage({
        role: Role.system,
        content: "System message",
      });
      const userMessage = new TextMessage({
        role: Role.user,
        content: "User message",
      });

      await adapter.process({
        threadId: "test-thread",
        model: "gpt-4o",
        messages: [systemMessage, userMessage],
        actions: [],
        eventSource: mockEventSource,
        forwardedParameters: {},
      });

      // Verify the stream function was called
      expect(mockEventSource.stream).toHaveBeenCalled();
    });

    it("should return the provided threadId", async () => {
      // Create a message
      const systemMessage = new TextMessage({
        role: Role.system,
        content: "System message",
      });

      const result = await adapter.process({
        threadId: "test-thread",
        model: "gpt-4o",
        messages: [systemMessage],
        actions: [],
        eventSource: mockEventSource,
        forwardedParameters: {},
      });

      expect(result.threadId).toBe("test-thread");
    });
  });
});