Class MockEmbeddingModel

java.lang.Object
dev.langchain4j.test.mock.MockEmbeddingModel
All Implemented Interfaces:
EmbeddingModel
@Experimental public class MockEmbeddingModel extends Object implements EmbeddingModel
An EmbeddingModel for tests: it needs no network, returns the vectors the test tells it to, and records the requests it receives.

Script the answers step by step, one step per request. One call of embed(...), embedAll(...) or embedAsync(...) is one request, also when it embeds several texts. A call of embedAll(...) without texts returns at once, without asking the model: it takes no step and records no request. For example, a model that returns the same vector for every text of every request: 

MockEmbeddingModel model = MockEmbeddingModel.thatEmbeds(0.1f, 0.2f).thenRepeats();
Every that... method creates a mock with its first step, 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 texts of 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. embed(...) throws it, and embedAsync(...) returns a future that fails with it. So a mock that is called more than once, for example to embed the documents of a search and later its query, ends its script with thenRepeats(), which answers every further request with the last step.

A script can let a request fail and the next one succeed, for example to test your own retry logic after a rate limit: 

MockEmbeddingModel model = MockEmbeddingModel.thatFails(new RateLimitException("Too many requests"))
        .thenEmbeds(0.1f, 0.2f)
        .thenRepeats();
or, when your code ingests documents in batches, with one ingest(...) or embedAll(...) call per batch, let its second batch fail and the next batches succeed: 
MockEmbeddingModel model = MockEmbeddingModel.thatEmbeds(0.1f, 0.2f)
        .thenFails(new InternalServerException("Service unavailable"))
        .thenEmbeds(0.1f, 0.2f)
        .thenRepeats();
A real model that splits one call into several requests to its provider fails the whole call when one of them fails, so one call is always one step.

Use thatEmbeds(Function) when the test depends on which texts are close to each other, for example on the ranking of retrieved segments or on a minimum score: the function in the test then says which texts are close. Use thatEmbeds(float...) when it does not matter. When a test depends on the meaning of texts, use a real embedding model that runs in-process, for example AllMiniLmL6V2QuantizedEmbeddingModel from the module langchain4j-embeddings-all-minilm-l6-v2-q.

Like on every mock in this package, thatResponds(...) and thenResponds(...) give a step as the whole response, here an EmbeddingResponse: the request is answered with exactly that response, including its metadata, for example the token usage of an embedding call in a cost test. The response is used as given, also when its number of embeddings differs from the number of texts of the request, for example to test how your code handles a provider that returns too few embeddings.

The mock accepts the per-call parameters modelName and inputType, and rejects dimensions with an UnsupportedFeatureException, like a model without that parameter; a rejected request takes no step. Unless a step gives the whole response, a response reports the model name of the request. embedAsync(...) completes on another thread, where a step that computes its vectors also computes them.

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

    • thenEmbeds

      public MockEmbeddingModel thenEmbeds(float... vector)
      Adds a step: every text of the next request gets the given vector.
      Throws:
      IllegalArgumentException - if the vector has no dimension
      IllegalStateException - if the script already ends with thenRepeats()

    • thenEmbeds

      public MockEmbeddingModel thenEmbeds(Function<String,float[]> embeddings)
      Adds a step: every text of the next request gets the vector computed from it, see thatEmbeds(Function).
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • thenResponds

      public MockEmbeddingModel thenResponds(EmbeddingResponse response)
      Adds a step: the next request is answered with exactly the given response, including its metadata, see thatResponds(EmbeddingResponse).
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • thenFails

      public MockEmbeddingModel 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 MockEmbeddingModel 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 vectors computes them again for each request. Without it, a request after the last step fails. It must be the last call of the script.
      Throws:
      IllegalStateException - if the script already ends with thenRepeats()

    • withListeners

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

    • withListeners

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

    • supportedParameters

      public Set<EmbeddingParameter<?>> supportedParameters()
      Description copied from interface: EmbeddingModel
      The per-call parameters this model honors. A request that populates a parameter outside this set is rejected by EmbeddingModel.embed(EmbeddingRequest) with an UnsupportedFeatureException, so a parameter a model cannot apply is never silently ignored. Defaults to an empty set (the model accepts no per-call parameters).
      Specified by:
      supportedParameters in interface EmbeddingModel

    • doEmbedAsync

      public CompletableFuture<EmbeddingResponse> doEmbedAsync(EmbeddingRequest request)
      Answers like doEmbed(EmbeddingRequest), with the next step of the script, and completes the future on another thread, where a step that computes its vectors also computes them. It never throws: a null request gives a failed future.
      Specified by:
      doEmbedAsync in interface EmbeddingModel
      Parameters:
      request - the request, with parameters already merged and validated.
      Returns:
      a CompletableFuture of the response.

    • doEmbed

      public EmbeddingResponse doEmbed(EmbeddingRequest request)
      Description copied from interface: EmbeddingModel
      Performs the embedding for EmbeddingModel.embed(EmbeddingRequest). This is the low-level method a provider overrides to build the actual API call; EmbeddingModel.embed(EmbeddingRequest) handles per-call parameter validation and listener dispatch around it.

      An implementation must override either this method (preferred) or EmbeddingModel.embedAll(List) — the two have mutually-recursive defaults, so overriding neither leads to infinite recursion. The default implementation embeds the text of each input via EmbeddingModel.embedAll(List), which keeps implementations that only provide EmbeddingModel.embedAll(List) working.

      When this method is called, the request's parameters have already been merged with EmbeddingModel.defaultRequestParameters() and validated against EmbeddingModel.supportedParameters() and EmbeddingModel.supportedContentTypes().

      Specified by:
      doEmbed in interface EmbeddingModel
      Parameters:
      request - the request, with parameters already merged and validated.
      Returns:
      the response.

    • dimension

      public int dimension()
      Embeds the text "test" with the first step of the script that returns vectors, and returns the length of its vector; a step that gives the whole response gives the length of its first embedding. It takes no step and records no request, so the script and requests() stay as they are. A script without such a step throws the exception of its first step, if that step fails, or an IllegalStateException. When an embedding store asks the model for its dimension, for example while a framework starts, set the dimension of the store explicitly instead.
      Specified by:
      dimension in interface EmbeddingModel
      Returns:
      dimension of the embedding

    • listeners

      public List<EmbeddingModelListener> listeners()
      Description copied from interface: EmbeddingModel
      The EmbeddingModelListeners that are notified around EmbeddingModel.embed(EmbeddingRequest). Configure them on the model's builder (via listeners(...)). Defaults to an empty list.
      Specified by:
      listeners in interface EmbeddingModel

    • requests

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

    • request

      public EmbeddingRequest request()
      The only request received so far.
      Throws:
      IllegalStateException - if the mock received no request or several requests

    • lastRequest

      public EmbeddingRequest lastRequest()
      The last request received so far.
      Throws:
      IllegalStateException - if the mock received no request

    • thatEmbeds

      public static MockEmbeddingModel thatEmbeds(float... vector)
      A model that answers the first request with the given vector for every text. All texts are then equally close to each other, so a test with this model cannot depend on their ranking. Add the next steps with then... methods, or return this vector for every request: thatEmbeds(0.1f, 0.2f).thenRepeats().
      Throws:
      IllegalArgumentException - if the vector has no dimension

    • thatEmbeds

      public static MockEmbeddingModel thatEmbeds(Function<String,float[]> embeddings)
      A model that answers the first request with the vector computed from each text, so the test says which texts are close to each other. For example, with 
      MockEmbeddingModel model = MockEmbeddingModel.thatEmbeds(
                text -> text.contains("refund") ? new float[] {1, 0} : new float[] {0, 1})
        .thenRepeats();
      two texts that both contain "refund" get a relevance score of 1, and a text with "refund" and a text without it get a score of 0.5. Add the next steps with then... methods, or compute the vectors of every request with thenRepeats(), as above. If the function throws, the request fails with that exception, unchanged. If it returns no vector, the request fails with an IllegalStateException that shows the text.

    • thatResponds

      public static MockEmbeddingModel thatResponds(EmbeddingResponse response)
      A model that answers the first request with exactly the given response, including its metadata, for example its token usage or model name, to test code that reads them, such as a cost check: 
      MockEmbeddingModel model = MockEmbeddingModel.thatResponds(EmbeddingResponse.builder()
        .embeddings(List.of(Embedding.from(new float[] {0.1f, 0.2f})))
        .tokenUsage(new TokenUsage(12))
        .build());
      The response is used as given, also when its number of embeddings differs from the number of texts of the request. Add the next steps with then... methods, or answer every request with this response: thatResponds(response).thenRepeats().

    • thatFails

      public static MockEmbeddingModel 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.