LLMs mit Informationen anreichern – RAG mit Spring AI

Die Serie „AI in Java“ geht weiter mit Retrieval Augmented Generation (RAG) in Spring AI. Nachdem wir uns im letzten Teil der Serie bereits mit diesem Thema und einer Umsetzung mithilfe von LangChain4j in Quarkus beschäftigt haben schauen wir nun eine weitere Umsetzung an.

Das Ziel ist wieder das gleiche: ein einfaches Beispiel in dem ein LLM (mit Ollama) auf Anfragen zu Verkehrs- und Wetterinformationen antwortet und die Datengrundlage für diese Antworten aus unseren bereitgestellten Dokumenten stammt. Wir nehmen diesmal das Modell llama3.2:3b, das etwas kleiner als das bisher verwendete mistral ist. Im Vergleich zu dem vorherigen Beispiel mit Easy RAG ist es hier etwas komplizierter initial in das Thema einzusteigen und es reicht nicht einfach eine Dependency einzubinden, eine Property zu definieren und erstmal mit der Standardkonfiguration zu starten. Mit ein klein wenig Aufwand ist es aber auch mit Spring AI einfach möglich ein simples RAG Beispielprojekt mit einem einfachen In-Memory Speicher zu erstellen.

Den Speicher müssen wir zunächst als Bean bereitstellen, wir können hier den SimpleVectorStore von Spring AI verwenden. Um das Befüllen müssen wir uns allerdings auch noch kümmern. Der Ordner mit den Dokumenten muss sich in src/main/resources befinden, um den Ordnernamen einfach konfigurieren zu können ist er in die Property de.gedoplan.showcase.springaidemo.rag.resourcepath ausgelagert. Von da können wir uns die einzelnen Dateipfade holen und eine Liste von Dokumenten erstellen. Auf geschachtelte Ordnerstrukturen sind wir in diesem simplen Beispiel noch nicht eingegangen.

@Configuration
public class VectorStoreConfig {
  @Bean
  VectorStore vectorStore(EmbeddingModel embeddingModel) {
    return SimpleVectorStore.builder(embeddingModel).build();
  }

  @Bean
  CommandLineRunner load(VectorStore vectorStore,
                         @Value("classpath:${de.gedoplan.showcase.springaidemo.rag.resourcepath}") String resourcePath) throws IOException {
    File ragResourceDir = new DefaultResourceLoader().getResource(resourcePath).getFile();
    if (!ragResourceDir.exists() && !ragResourceDir.isDirectory()) {
      throw new IllegalArgumentException("Resource path '" + resourcePath + "' does not exist or is not a directory");
    }
    File[] files = ragResourceDir.listFiles();
    if (files == null || files.length == 0) {
      throw new IllegalArgumentException("No files found in resource path '" + resourcePath + "'");
    }
    List<Document> documents = new ArrayList<>();
    Arrays.stream(files).forEach(file -> {
      if (!file.isDirectory()) {
        documents.addAll(
          new TokenTextSplitter().transform(
            new TextReader(resourcePath + File.separator + file.getName()).get()
          )
        );
      } else {
        throw new IllegalArgumentException("Nested RAG resource directories not supported yet.");
      }
    });
    return ignored -> vectorStore.add(documents);
  }
}

Die Daten werden so bei jedem Start der Anwendung neu eingelesen und verarbeitet. Denkbar wäre hier auch eine Implementierung, die die Embeddings in einer JSON-Datei abspeichert und diese dann bei einem Neustart nicht neu bestimmen muss, ähnlich wie es bei Easy RAG in LangChain4j konfigurierbar ist.

Auch in diesem Beispielprojekt ist eine HTTP-Schnittstelle implementiert um einen kleinen Überblick über die Embedding-Scores zu bekommen.

@PostMapping(consumes = MediaType.TEXT_PLAIN_VALUE, produces = MediaType.TEXT_PLAIN_VALUE, value = "/checkEmbeddingScore")
public String checkEmbeddingScore(@RequestBody String question) {
  List<Document> results = this.vectorStore.similaritySearch(SearchRequest.builder().similarityThresholdAll().topK(10).query(question).build());
  return results.stream().map(document -> {
    Double score = document.getScore();
    return score + ": " + document.getContent();
  }).collect(Collectors.joining("\n\n"));
}

What is the current weather in Bielefeld?

0.4504619925388119: Weather forecast for Berlin:

Today: 25°C/14°C, sunny
Tomorrow: 26°C/15°C, partly cloudy
Day after tomorrow: 24°C/14°C, mostly sunny

0.44789886725109695: Streetworks on the A2:

From March 15, 2024 to August 24, 2024, bridge construction work between Bielefeld and Hanover.
From January 23, 2024 to February 24, 2024, two-lane traffic between Hanover and Oberhausen.
From June 1, 2024 to June 10, 2024, full closure near Dortmund.

0.4401144314803378: Streetworks on the A1:

From July 15, 2024 to December 24, 2024, bridge construction work between Cologne and Dortmund.
From July 23, 2024 to July 24, 2024, single-lane traffic between Osnabrück and Bremen.
From January 1, 2024 to January 10, 2024, full closure near Hamburg.

0.40504789201978736: Weather forecast for Bielefeld:

Today: 19°C/12°C, partly cloudy
Tomorrow: 22°C/13°C, mostly sunny
Day after tomorrow: 22°C/12°C, partly cloudy

Wenn wir uns die Frage anschauen, dann hätten wir hier vermutlich erwartet, dass die relevanteste Information die Wettervorhersage für Bielefeld sein müsste. Diese hat aber nur einen Score von 0,4 und liegt damit z.B. vom Score her hinter der Wettervorhersage für Berlin. Spring AI verwendet mit Ollama standardmäßig mistral als AI Modell um die Embeddings zu bestimmen. Sinnvoller kann es hier sein ein auf Embeddings spezialisiertes AI Modell zu verwenden. Über eine Property können wir das Modell umstellen und zum Beispiel dasselbe Modell nutzen, das LangChain4j standardmäßig nutzt: spring.ai.ollama.embedding.options.model=nomic-embed-text:latest. Damit ist das initiale Befüllen der Datenbank deutlich schneller und die Scores entsprechen auch etwas eher unseren Erwartungen.

What is the current weather in Bielefeld?

0.8523263139092256: Weather forecast for Bielefeld:

Today: 19°C/12°C, partly cloudy
Tomorrow: 22°C/13°C, mostly sunny
Day after tomorrow: 22°C/12°C, partly cloudy

0.7115280685411207: Weather forecast for Berlin:

Today: 25°C/14°C, sunny
Tomorrow: 26°C/15°C, partly cloudy
Day after tomorrow: 24°C/14°C, mostly sunny

0.6971158909686656: Weather forecast for Cologne:

Today: 22°C/13°C, sunny
Tomorrow: 24°C/14°C, sunny
Day after tomorrow: 24°C/12°C, sunny

0.5962977983343077: Streetworks on the A2:

From March 15, 2024 to August 24, 2024, bridge construction work between Bielefeld and Hanover.
From January 23, 2024 to February 24, 2024, two-lane traffic between Hanover and Oberhausen.
From June 1, 2024 to June 10, 2024, full closure near Dortmund.

Die Kontextinformationen aus unseren Dokumenten können wir einer Anfrage an den ChatClient mit einem QuestionAnswerAdvisor hinzufügen. Diesen können wir entweder einfach nur mit dem VectorStore erstellen oder auch weiter konfigurieren. So können wir z.B. das Prompt-Template anpassen oder Suchparameter wie eine untere Grenze für den Score oder eine maximale Anzahl an Ergebnissen einstellen. Wichtig zu beachten ist hierbei, dass der Score für die Ähnlichkeit hier einen Wert zwischen -1 und 1 haben kann. Die Grenze von 0,8, die wir in dem Beispielprojekt mit LangChain4j gewählt haben, entspricht hier also einer Grenze von 0,6.

@PostMapping(consumes = MediaType.TEXT_PLAIN_VALUE, produces = MediaType.TEXT_PLAIN_VALUE)
public String rag(@RequestBody String question) {
  String RAG_PROMPT = """
          
    To answer this question you can use the following information if applicable:
    {question_answer_context}\
    """;
  return chatClient
    .prompt()
    .user(question)
    .advisors(new QuestionAnswerAdvisor(
      vectorStore,
      SearchRequest.builder().similarityThreshold(0.6).topK(3).build(),
      RAG_PROMPT))
    .call()
    .content();
}

Das Beispielprojekt gibt es wie immer auch auf GitHub.