Class MockChatModel

java.lang.Object
dev.langchain4j.test.mock.MockChatModel
All Implemented Interfaces:
ChatModel
@Experimental public class MockChatModel extends Object implements ChatModel
A ChatModel for tests: it needs no network, answers the way the test tells it to, and records the requests it receives.

Script the answers step by step, one step per request. For example, a model that first calls a tool and then answers with text: 

MockChatModel model = MockChatModel.thatCallsTool("cancelBooking", Map.of("bookingNumber", "MS-777"))
        .thenResponds("Booking MS-777 is cancelled.");

Assistant assistant = AiServices.builder(Assistant.class)
        .chatModel(model)
        .tools(new BookingTools())
        .build();
Every that... method creates a mock with its first step (one step per message for the message-list forms), and every then... method adds a step. Each step answers exactly one request. A request after the last step fails with an IllegalStateException that shows the request, so a test notices when the code calls the model more often than the script expects, unless the code under test handles the error itself. chat(...) throws it, and chatAsync(...) returns a future that fails with it.

To answer every further request with the last step, end the script with thenRepeats(). For example, a model that gives the same answer to every request: 

MockChatModel model = MockChatModel.thatResponds("Hello!").thenRepeats();
A step can compute its answer from the request with thatResponds(Function), for example to answer concurrent requests by their content: thatResponds(request -> ...).thenRepeats().

To check what the model received, use request() or requests(), for example with LangChain4jAssertions.assertThat(model.request()).hasLastUserMessageContaining("MS-777"). To check what an AI Service did with the answers, capture the invocation with AiServiceCapture.

A response reports what a provider reports: the model name of the request, and the finish reason FinishReason.TOOL_EXECUTION for an answer with tool calls, otherwise FinishReason.STOP. To report token usage in every response, use withTokenUsage(TokenUsage). To report anything else, for example FinishReason.LENGTH or a different token usage per step, give the whole response: thatResponds(ChatResponse).

By default, the mock reports no capabilities, so an AI Service that returns a structured type, for example a POJO, adds format instructions to the user message. To make the AI Service ask for a JSON schema response format instead, as it does with a model that supports it, use withSupportedCapabilities(Capability.RESPONSE_FORMAT_JSON_SCHEMA).

A mock remembers which steps were taken and which requests it received, so create a new one for each test. A mock shared between tests, for example as a bean, fails with a clear message as soon as its script runs out, unless the script ends with thenRepeats(). The steps are taken in the order the requests arrive. A mock keeps every request it receives, so it is not meant for load tests or other long runs.

  • Method Details

    • thenResponds

      public MockChatModel thenResponds(String text)
      Adds a step: the next request is answered with the given text.
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • thenResponds

      public MockChatModel thenResponds(AiMessage aiMessage)
      Adds a step: the next request is answered with the given message.
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • thenResponds

      public MockChatModel thenResponds(ChatResponse chatResponse)
      Adds a step: the next request is answered with exactly the given response, including its metadata, for example its token usage or finish reason. Only tool calls without an id get the next id of the mock, see toolCall(String, Map).
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • thenCallsTool

      public MockChatModel thenCallsTool(String toolName, Map<String,?> arguments)
      Adds a step: the next request is answered with a call of the given tool.
      Parameters:
      toolName - the name of the tool
      arguments - the arguments of the tool, by name
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • thenCallsTool

      public MockChatModel thenCallsTool(String toolName, String jsonArguments)
      Adds a step: the next request is answered with a call of the given tool. Use it for arguments that cannot be given as a map, for example invalid JSON.
      Parameters:
      toolName - the name of the tool
      jsonArguments - the arguments of the tool, as JSON
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • thenFails

      public MockChatModel thenFails(RuntimeException exception)
      Adds a step: the next request fails with exactly the given exception.
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • thenRepeats

      public MockChatModel thenRepeats()
      Answers every further request with the last step of the script. Only the last step repeats; the script does not start over. A last step that computes its answer computes it again for each request. Without it, a request after the last step fails. It must be the last call of the script.

      A step that calls a tool cannot repeat. To test a model that keeps calling a tool, for example the limit set with maxToolCallingRoundTrips, compute the answer: thatResponds(request -> AiMessage.from(toolCall("getWeather", Map.of()))).thenRepeats().

      Throws:
      IllegalStateException - if the last step calls a tool, because an AI Service would then call the tool again and again, or if the script already ends with thenRepeats()

    • withTokenUsage

      public MockChatModel withTokenUsage(@Nullable TokenUsage tokenUsage)
      Reports the given token usage in the response to every request, except for a step that gives the whole response: that one reports exactly its own token usage. For different token usages per step, give the whole responses, see thenResponds(ChatResponse).
      Parameters:
      tokenUsage - the token usage, or null for no token usage, which is also the default

    • withListeners

      public MockChatModel withListeners(ChatModelListener... listeners)
      Notifies the given listeners of each call, like a real model does.

    • withListeners

      public MockChatModel withListeners(List<ChatModelListener> listeners)
      Notifies the given listeners of each call, like a real model does.

    • withSupportedCapabilities

      public MockChatModel withSupportedCapabilities(Capability... capabilities)
      Reports the given capabilities from supportedCapabilities(), like a model that supports them. For example, with Capability.RESPONSE_FORMAT_JSON_SCHEMA, an AI Service that returns a structured type, for example a POJO, asks for a JSON schema response format instead of adding format instructions to the user message. The mock only reports the capabilities: it still answers with the steps of its script, and it does not check the requests or the answers against the capabilities.

    • withSupportedCapabilities

      public MockChatModel withSupportedCapabilities(Set<Capability> capabilities)
      Reports the given capabilities from supportedCapabilities(), see withSupportedCapabilities(Capability...).

    • listeners

      public List<ChatModelListener> listeners()
      Specified by:
      listeners in interface ChatModel

    • supportedCapabilities

      public Set<Capability> supportedCapabilities()
      Specified by:
      supportedCapabilities in interface ChatModel

    • doChat

      public ChatResponse doChat(ChatRequest chatRequest)
      Specified by:
      doChat in interface ChatModel

    • doChatAsync

      public CompletableFuture<ChatResponse> doChatAsync(ChatRequest chatRequest)
      Answers like doChat(ChatRequest), with the next step of the script, and completes the future on another thread, where a step that computes its answer also computes it. It does not call doChat(ChatRequest).
      Specified by:
      doChatAsync in interface ChatModel
      Parameters:
      chatRequest - a ChatRequest, containing all the inputs to the LLM
      Returns:
      a CompletableFuture of the ChatResponse

    • requests

      public List<ChatRequest> requests()
      A copy of the requests received so far, in the order they arrived.

    • request

      public ChatRequest request()
      The only request received so far. To check it, for example: LangChain4jAssertions.assertThat(model.request()).hasLastUserMessageContaining("MS-777").
      Throws:
      IllegalStateException - if the mock received no request or several requests

    • lastRequest

      public ChatRequest lastRequest()
      The last request received so far. For example, the request that follows the tool results.
      Throws:
      IllegalStateException - if the mock received no request

    • thatResponds

      public static MockChatModel thatResponds(String text)
      A model that answers the first request with the given text. Add the next steps with then... methods, or answer every request with this text: thatResponds("Hello!").thenRepeats().

    • thatResponds

      public static MockChatModel thatResponds(AiMessage... aiMessages)
      A model that answers with the given messages in order, one per request. Add further steps with then... methods. Use it for answers that thatCallsTool(String, Map) cannot express, for example several tool calls in one answer: AiMessage.from(toolCall("a", argsA), toolCall("b", argsB)).
      Throws:
      IllegalArgumentException - if no message is given

    • thatResponds

      public static MockChatModel thatResponds(List<AiMessage> aiMessages)
      A model that answers with the given messages in order, one per request. Add further steps with then... methods.
      Throws:
      IllegalArgumentException - if the list is empty

    • thatResponds

      public static MockChatModel thatResponds(ChatResponse chatResponse)
      A model that answers the first request with exactly the given response, including its metadata, for example its token usage, model name or finish reason. Only tool calls without an id get the next id of the mock, see toolCall(String, Map). Add the next steps with then... methods.

    • thatResponds

      public static MockChatModel thatResponds(Function<ChatRequest, AiMessage> answers)
      A model that computes the answer to the first request from the request. Add the next steps with then... methods, or compute the answer to every request with thenRepeats(). For example, a model that answers by the user's question: 
      MockChatModel model = MockChatModel.thatResponds(request -> {
    String question = UserMessage.findLast(request.messages()).orElseThrow().singleText();
    return AiMessage.from(question.contains("refund") ? "Refunds take 5 days." : "How can I help?");
}).thenRepeats();
      If the function throws, the request fails with that exception, unchanged. With chatAsync(...), the function runs on another thread.

    • thatCallsTool

      public static MockChatModel thatCallsTool(String toolName, Map<String,?> arguments)
      A model that answers the first request with a call of the given tool. Add the next steps with then... methods. The call gets the next id of the mock, see toolCall(String, Map).
      Parameters:
      toolName - the name of the tool
      arguments - the arguments of the tool, by name

    • thatCallsTool

      public static MockChatModel thatCallsTool(String toolName, String jsonArguments)
      A model that answers the first request with a call of the given tool. Add the next steps with then... methods. Use it for arguments that cannot be given as a map, for example invalid JSON.
      Parameters:
      toolName - the name of the tool
      jsonArguments - the arguments of the tool, as JSON

    • thatFails

      public static MockChatModel thatFails(RuntimeException exception)
      A model whose first request fails with exactly the given exception. To let every request fail, add thenRepeats(). To let later requests succeed, for example to test your own retry logic, add steps with then... methods.

    • toolCall

      public static ToolExecutionRequest toolCall(String toolName, Map<String,?> arguments)
      A call of the given tool without an id, to build an answer with AiMessage.from(...), for example in thatResponds(Function) or for several tool calls in one answer. When a mock answers with it, the call gets the next id of that mock: tool_call_1, tool_call_2, and so on, so that every call has its own id, also two calls of the same tool with the same arguments. To give a call a fixed id, use ToolExecutionRequest.builder(). An answer with two calls that have the same fixed id is rejected, because a provider never gives two calls of one answer the same id. The ids are unique within one mock only: when several mocks answer into one chat memory, for example two AI Services that share it, give the calls fixed ids.
      Parameters:
      toolName - the name of the tool
      arguments - the arguments of the tool, by name

    • toolCall

      public static ToolExecutionRequest toolCall(String toolName, String jsonArguments)
      A call of the given tool without an id, like toolCall(String, Map), for arguments that cannot be given as a map, for example invalid JSON.
      Parameters:
      toolName - the name of the tool
      jsonArguments - the arguments of the tool, as JSON