Blog

LangChain4j com Elasticsearch como armazenamento de incorporação

O LangChain4j (LangChain para Java) utiliza o Elasticsearch como armazenamento integrado. Descubra como usá-lo para construir sua aplicação RAG em Java puro.

Na publicação anterior, descobrimos o que é LangChain4j e como:

  • Inicie uma discussão com os LLMs implementando um ChatLanguageModel e um ChatMemory

  • Manter o histórico do chat na memória para relembrar o contexto de uma discussão anterior com um LLM

Esta postagem do blog aborda como:

  • Criar vetores incorporados a partir de exemplos de texto

  • Armazene os vetores de incorporação no repositório de incorporações do Elasticsearch.

  • Buscar vetores semelhantes

Criar incorporações

Para criar embeddings, precisamos definir um EmbeddingModel para usar. Por exemplo, podemos usar o mesmo modelo mistral que usamos na postagem anterior. Estava correndo com a lhama:

EmbeddingModel model = OllamaEmbeddingModel.builder()
  .baseUrl(ollama.getEndpoint())
  .modelName(MODEL_NAME)
  .build();

Um modelo é capaz de gerar vetores a partir de texto. Aqui podemos verificar o número de dimensões geradas pelo modelo:

Logger.info("Embedding model has {} dimensions.", model.dimension());
// This gives: Embedding model has 4096 dimensions.

Para gerar vetores a partir de um texto, podemos usar:

Response<Embedding> response = model.embed("A text here");

Ou, se também quisermos fornecer metadados para nos permitir filtrar por coisas como texto, preço, data de lançamento ou qualquer outra coisa, podemos usar Metadata.from(). Por exemplo, estamos adicionando aqui o nome do jogo como um campo de metadados:

TextSegment game1 = TextSegment.from("""
    The game starts off with the main character Guybrush Threepwood stating "I want to be a pirate!"
    To do so, he must prove himself to three old pirate captains. During the perilous pirate trials, 
    he meets the beautiful governor Elaine Marley, with whom he falls in love, unaware that the ghost pirate 
    LeChuck also has his eyes on her. When Elaine is kidnapped, Guybrush procures crew and ship to track 
    LeChuck down, defeat him and rescue his love.
""", Metadata.from("gameName", "The Secret of Monkey Island"));
Response<Embedding> response1 = model.embed(game1);
TextSegment game2 = TextSegment.from("""
    Out Run is a pseudo-3D driving video game in which the player controls a Ferrari Testarossa 
    convertible from a third-person rear perspective. The camera is placed near the ground, simulating 
    a Ferrari driver's position and limiting the player's view into the distance. The road curves, 
    crests, and dips, which increases the challenge by obscuring upcoming obstacles such as traffic 
    that the player must avoid. The object of the game is to reach the finish line against a timer.
    The game world is divided into multiple stages that each end in a checkpoint, and reaching the end 
    of a stage provides more time. Near the end of each stage, the track forks to give the player a 
    choice of routes leading to five final destinations. The destinations represent different 
    difficulty levels and each conclude with their own ending scene, among them the Ferrari breaking 
    down or being presented a trophy.
""", Metadata.from("gameName", "Out Run"));
Response<Embedding> response2 = model.embed(game2);

Se você quiser executar este código, consulte a classe Step5EmbedddingsTest.java .

Adicionar o Elasticsearch para armazenar nossos vetores.

LangChain4j fornece um armazenamento de incorporação em memória. Isso é útil para executar testes simples:

EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
embeddingStore.add(response1.content(), game1);
embeddingStore.add(response2.content(), game2);

Mas, obviamente, isso não funcionaria com conjuntos de dados muito maiores, porque esse armazenamento de dados guarda tudo na memória e não temos memória infinita em nossos servidores. Assim, poderíamos armazenar nossos embeddings no Elasticsearch, que é, por definição, "elástico" e pode ser dimensionado verticalmente e horizontalmente conforme a demanda de dados. Para isso, vamos adicionar o Elasticsearch ao nosso projeto:

<dependency>
  <groupId>dev.langchain4j</groupId>
  <artifactId>langchain4j-elasticsearch</artifactId>
  <version>${langchain4j.version}</version>
</dependency>

<dependency>
  <groupId>org.testcontainers</groupId>
  <artifactId>elasticsearch</artifactId>
  <version>1.20.1</version>
  <scope>test</scope>
</dependency>

Como você deve ter notado, também adicionamos o módulo Elasticsearch TestContainers ao projeto, para que possamos iniciar uma instância do Elasticsearch a partir de nossos testes:

// Create the elasticsearch container
ElasticsearchContainer container =
  new ElasticsearchContainer("docker.elastic.co/elasticsearch/elasticsearch:8.15.0")
    .withPassword("changeme");

// Start the container. This step might take some time...
container.start();

// As we don't want to make our TestContainers code more complex than
// needed, we will use login / password for authentication.
// But note that you can also use API keys which is preferred.
final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials("elastic", "changeme"));

// Create a low level Rest client which connects to the elasticsearch container.
client = RestClient.builder(HttpHost.create("https://" + container.getHttpHostAddress()))
  .setHttpClientConfigCallback(httpClientBuilder -> {
    httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
    httpClientBuilder.setSSLContext(container.createSslContextFromCa());
    return httpClientBuilder;
  })
  .build();

// Check the cluster is running
client.performRequest(new Request("GET", "/"));

Para usar o Elasticsearch como um armazenamento de dados incorporado, você "simplesmente" precisa trocar o armazenamento de dados em memória do LangChain4j pelo armazenamento de dados do Elasticsearch:

EmbeddingStore<TextSegment> embeddingStore =
  ElasticsearchEmbeddingStore.builder()
    .restClient(client)
    .build();
embeddingStore.add(response1.content(), game1);
embeddingStore.add(response2.content(), game2);

Isso armazenará seus vetores no Elasticsearch em um índice default . Você também pode alterar o nome do índice para algo mais significativo:

EmbeddingStore<TextSegment> embeddingStore =
  ElasticsearchEmbeddingStore.builder()
    .indexName("games")
    .restClient(client)
    .build();
embeddingStore.add(response1.content(), game1);
embeddingStore.add(response2.content(), game2);

Se você quiser executar este código, consulte a classe Step6ElasticsearchEmbedddingsTest.java .

Buscar vetores semelhantes

Para buscar vetores semelhantes, primeiro precisamos transformar nossa pergunta em uma representação vetorial usando o mesmo modelo que usamos anteriormente. Já fizemos isso, então não é difícil fazer de novo. Note que, neste caso, não precisamos dos metadados:

String question = "I want to pilot a car";
Embedding questionAsVector = model.embed(question).content();

Podemos construir uma solicitação de pesquisa com essa representação da nossa pergunta e pedir ao repositório de embeddings para encontrar os primeiros vetores principais:

EmbeddingSearchResult<TextSegment> result = embeddingStore.search(
  EmbeddingSearchRequest.builder()
    .queryEmbedding(questionAsVector)
    .build());

Agora podemos iterar sobre os resultados e imprimir algumas informações, como o nome do jogo, que vem dos metadados, e a pontuação:

result.matches().forEach(m -> Logger.info("{} - score [{}]",
  m.embedded().metadata().getString("gameName"), m.score()));

Como era de se esperar, isso nos dá "Out Run" como o primeiro sucesso:

Out Run
Out Run - score [0.86672974]
The Secret of Monkey Island - score [0.85569763]

Se você quiser executar este código, consulte a classe Step7SearchForVectorsTest.java .

Nos bastidores

A configuração padrão do Elasticsearch Embedding Store utiliza a consulta kNN aproximada nos bastidores.

POST games/_search
{
  "query" : {
    "knn": {
      "field": "vector",
      "query_vector": [-0.019137882, /* ... */, -0.0148779955]
    }
  }
}

Mas isso poderia ser alterado fornecendo uma configuração diferente (ElasticsearchConfigurationScript) da configuração padrão (ElasticsearchConfigurationKnn) para o armazenamento de incorporação:

EmbeddingStore<TextSegment> embeddingStore =
  ElasticsearchEmbeddingStore.builder()
    .configuration(ElasticsearchConfigurationScript.builder().build())
    .indexName("games")
    .restClient(client)
    .build();

A implementação ElasticsearchConfigurationScript executa em segundo plano uma consulta script_score usando uma função cosineSimilarity .

Basicamente, ao ligar:

EmbeddingSearchResult<TextSegment> result = embeddingStore.search(
  EmbeddingSearchRequest.builder()
    .queryEmbedding(questionAsVector)
    .build());

Isto agora exige:

POST games/_search
{
  "query": {
    "script_score": {
      "script": {
        "source": "(cosineSimilarity(params.query_vector, 'vector') + 1.0) / 2",
        "params": {
          "queryVector": [-0.019137882, /* ... */, -0.0148779955]
        }
      }
    }
  }
}

Nesse caso, o resultado não muda em termos de "ordem", apenas a pontuação é ajustada, pois a chamada cosineSimilarity não usa nenhuma aproximação, mas calcula o cosseno para cada um dos vetores correspondentes:

Out Run - score [0.871952]
The Secret of Monkey Island - score [0.86380446]

Se você quiser executar este código, consulte a classe Step7SearchForVectorsTest.java .

Conclusão

Já abordamos a facilidade com que você pode gerar embeddings a partir do seu texto e como você pode armazenar e pesquisar os vizinhos mais próximos no Elasticsearch usando duas abordagens diferentes:

  • Usando a consulta aproximada e rápida knn com a opção padrão ElasticsearchConfigurationKnn

  • Usando a consulta exata, porém mais lenta, script_score com a opção ElasticsearchConfigurationScript

O próximo passo será construir uma aplicação RAG completa, com base no que aprendemos aqui.

Conteúdo relacionado

Testando seu código Java com mocks e Elasticsearch real.

Piotr Przybyl

Apresentamos o LangChain4j para simplificar a integração do LLM em aplicações Java.

David Pilato

Pronto para criar buscas de última geração?

Uma pesquisa suficientemente avançada não se consegue apenas com o esforço de uma só pessoa. O Elasticsearch é impulsionado por cientistas de dados, especialistas em operações de aprendizado de máquina, engenheiros e muitos outros que são tão apaixonados por buscas quanto você. Vamos nos conectar e trabalhar juntos para construir a experiência de busca mágica que lhe trará os resultados desejados.

Experimente você mesmo(a)