Blog

LangChain4j con Elasticsearch como almacén de incrustaciones

LangChain4j (LangChain para Java) tiene Elasticsearch como almacén de incrustación. Descubre cómo usarlo para construir tu aplicación RAG en Java puro.

En la publicación anterior, descubrimos qué es LangChain4j y cómo:

  • Habla con los LLMs implementando un ChatLanguageModel y un ChatMemory

  • Conserva el historial de chats en memoria para recordar el contexto de una conversación previa con un LLM

Esta entrada del blog trata sobre cómo:

  • Crear incrustaciones vectoriales a partir de ejemplos de texto

  • Almacenar incrustaciones vectoriales en el almacenamiento de incrustaciones de Elasticsearch

  • Búsqueda de vectores similares

Crear incrustaciones

Para crear embeddings, necesitamos definir un EmbeddingModel a emplear. Por ejemplo, podemos usar el mismo modelo de mistral que usamos en la publicación anterior. Estaba corriendo con ollama:

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

Un modelo es capaz de generar vectores a partir de texto. Aquí podemos comprobar el número de dimensiones generadas por el modelo:

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

Para generar vectores a partir de un texto, podemos usar:

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

O si también queremos proporcionar metadatos para filtrar cosas como texto, precio, fecha de lanzamiento o lo que sea, podemos usar Metadata.from(). Por ejemplo, aquí agregamos el nombre del juego como campo de metadatos:

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);

Si quieres ejecutar este código, por favor echa un vistazo a la clase Step5EmbedddingsTest.java .

Agregar Elasticsearch para almacenar nuestros vectores

LangChain4j proporciona un almacén de incrustación en memoria. Esto es útil para realizar pruebas sencillas:

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

Pero obviamente, esto no podría funcionar con un conjunto de datos mucho mayor porque este almacén almacena todo en memoria y no tenemos memoria infinita en nuestros servidores. Así que, en su lugar, podríamos almacenar nuestras incrustaciones en Elasticsearch, que por definición es "elástico" y puede escalar hacia arriba y hacia fuera con tus datos. Para ello, agregamos Elasticsearch a nuestro proyecto:

<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 notasteis, también agregamos el módulo Elasticsearch TestContainers al proyecto, para poder iniciar una instancia de Elasticsearch a partir de nuestras pruebas:

// 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 Elasticsearch como almacén de embebed, "solo" tienes que cambiar del almacén de datos en memoria de LangChain4j al de Elasticsearch:

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

Esto almacenará tus vectores en Elasticsearch en un índice default . También puedes cambiar el nombre del índice por algo más significativo:

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

Si quieres ejecutar este código, por favor echa un vistazo a la clase Step6ElasticsearchEmbedddingsTest.java .

Búsqueda de vectores similares

Para buscar vectores similares, primero necesitamos transformar nuestra pregunta en una representación vectorial usando el mismo modelo que usamos anteriormente. Ya lo hicimos, así que no es difícil hacerlo de nuevo. Ten en cuenta que en este caso no necesitamos los metadatos:

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

Podemos construir una solicitud de búsqueda con esta representación de nuestra pregunta y pedir al almacén de incrustaciones que encuentre los primeros vectores principales:

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

Podemos iterar sobre los resultados ahora e imprimir algo de información, como el nombre del juego que proviene de los metadatos y el puntaje:

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

Como era de esperar, esto nos da "Out Run" como primer éxito:

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

Si quieres ejecutar este código, por favor echa un vistazo a la clase Step7SearchForVectorsTest.java .

Tras bambalinas

La configuración predeterminada para la tienda de incrustación de Elasticsearch emplea la consulta aproximada kNN detrás de escena.

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

Pero esto podría cambiar proporcionando otra configuración (ElasticsearchConfigurationScript) distinta a la predeterminada (ElasticsearchConfigurationKnn) en la tienda de incrustación:

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

La implementación ElasticsearchConfigurationScript ejecuta entre bastidores una consulta script_score usando una función cosineSimilarity.

Básicamente, al llamar:

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

Esto ahora llama:

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

En ese caso, el resultado no cambia en términos de "orden", sino que solo se ajusta el puntaje porque la llamada cosineSimilarity no emplea ninguna aproximación, sino que calcula el coseno para cada uno de los vectores emparejados:

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

Si quieres ejecutar este código, por favor echa un vistazo a la clase Step7SearchForVectorsTest.java .

Conclusión

Explicamos lo fácil que es generar incrustaciones a partir de tu texto y cómo puedes almacenar y buscar los vecinos más cercanos en Elasticsearch usando dos enfoques diferentes:

  • Usando la consulta de knn aproximada y rápida con la opción de ElasticsearchConfigurationKnn por defecto

  • Usando la consulta de script_score exacta pero más lenta con la opción ElasticsearchConfigurationScript

El siguiente paso será construir una aplicación completa de RAG, basada en lo que aprendimos aquí.

Contenido relacionado

Probar tu código Java con mocks y Elasticsearch real

Piotr Przybyl

Presentamos LangChain4j para simplificar la integración de LLM en aplicaciones Java

David Pilato

¿Estás listo para crear experiencias de búsqueda de última generación?

No se logra una búsqueda suficientemente avanzada con los esfuerzos de uno. Elasticsearch está impulsado por científicos de datos, operaciones de ML, ingenieros y muchos más que son tan apasionados por la búsqueda como tú. Conectemos y trabajemos juntos para crear la experiencia mágica de búsqueda que te dará los resultados que deseas.

Pruébalo tú mismo