Enhance use cases with Ollama and PGVector

ThomasVitale · Jul 28, 2024 · 2ca40cd · 2ca40cd
1 parent 5d774da
commit 2ca40cd
Show file tree

Hide file tree

Showing 17 changed files with 87 additions and 49 deletions.
diff --git a/00-use-cases/question-answering/README.md b/00-use-cases/question-answering/README.md
@@ -1,37 +1,39 @@
 # Question Answering (RAG)
 
-Ask questions about documents with LLMs via Ollama.
+Ask questions about documents with LLMs via Ollama and PGVector.
 
 ## Running the application
 
 The application relies on Ollama for providing LLMs. You can either run Ollama locally on your laptop, or rely on the Testcontainers support in Spring Boot to spin up an Ollama service automatically.
+Furthermore, the application relies on the native Testcontainers support in Spring Boot to spin up a PostgreSQL database with the pgvector extension for embeddings.
 
 ### Ollama as a native application
 
 First, make sure you have [Ollama](https://ollama.ai) installed on your laptop.
-Then, use Ollama to run the _mistral_ large language model.
+Then, use Ollama to run the _mistral_ and _nomic-embed-text_ large language models.
 
 ```shell
 ollama run mistral
+ollama run nomic-embed-text
 ```
 
 Finally, run the Spring Boot application.
 
 ```shell
-./gradlew bootRun
+./gradlew bootTestRun
 ```
 
 ### Ollama as a dev service with Testcontainers
 
-The application relies on the native Testcontainers support in Spring Boot to spin up an Ollama service with a _mistral_ model at startup time.
+The application relies on the native Testcontainers support in Spring Boot to spin up an Ollama service with _mistral_ and _nomic-embed-text_ models at startup time.
 
 ```shell
-./gradlew bootTestRun
+./gradlew bootTestRun --args='--spring.profiles.active=ollama-image'
 ```
 
 ## Calling the application
 
-You can now call the application that will use Ollama and _mistral_ to load text documents as embeddings and generate an answer to your questions based on those documents (RAG pattern).
+You can now call the application that will use Ollama with _nomic-embed-text_ and _mistral_ to load text documents as embeddings and generate an answer to your questions based on those documents (RAG pattern).
 This example uses [httpie](https://httpie.io) to send HTTP requests.
 
 ```shell

diff --git a/00-use-cases/question-answering/build.gradle b/00-use-cases/question-answering/build.gradle
@@ -24,6 +24,7 @@ dependencies {
 
     implementation 'org.springframework.boot:spring-boot-starter-web'
     implementation 'org.springframework.ai:spring-ai-ollama-spring-boot-starter'
+    implementation 'org.springframework.ai:spring-ai-pgvector-store-spring-boot-starter'
 
     testAndDevelopmentOnly 'org.springframework.boot:spring-boot-devtools'
 
@@ -33,6 +34,7 @@ dependencies {
     testImplementation 'org.springframework.ai:spring-ai-spring-boot-testcontainers'
     testImplementation 'org.testcontainers:junit-jupiter'
     testImplementation 'org.testcontainers:ollama'
+    testImplementation 'org.testcontainers:postgresql'
 }
 
 tasks.named('test') {

diff --git a/00-use-cases/question-answering/src/main/java/com/thomasvitale/ai/spring/ChatService.java b/00-use-cases/question-answering/src/main/java/com/thomasvitale/ai/spring/ChatService.java
@@ -2,7 +2,6 @@
 
 import org.springframework.ai.chat.client.ChatClient;
 import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
-import org.springframework.ai.vectorstore.SearchRequest;
 import org.springframework.ai.vectorstore.VectorStore;
 import org.springframework.stereotype.Service;
 
@@ -19,7 +18,7 @@ class ChatService {
 
     String chatWithDocument(String message) {
         return chatClient.prompt()
-                .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults().withTopK(5)))
+                .advisors(new QuestionAnswerAdvisor(vectorStore))
                 .user(message)
                 .call()
                 .content();

diff --git a/...-cases/question-answering/src/main/java/com/thomasvitale/ai/spring/IngestionPipeline.java b/...-cases/question-answering/src/main/java/com/thomasvitale/ai/spring/IngestionPipeline.java
@@ -48,7 +48,7 @@ public void run() {
         documents.addAll(textReader2.get());
 
         logger.info("Creating and storing Embeddings from Documents");
-        vectorStore.add(new TokenTextSplitter(300, 300, 5, 1000, true).split(documents));
+        vectorStore.add(new TokenTextSplitter().split(documents));
     }
 
 }
diff --git a/...-cases/question-answering/src/main/java/com/thomasvitale/ai/spring/QuestionAnswering.java b/...-cases/question-answering/src/main/java/com/thomasvitale/ai/spring/QuestionAnswering.java
@@ -1,11 +1,7 @@
 package com.thomasvitale.ai.spring;
 
-import org.springframework.ai.embedding.EmbeddingModel;
-import org.springframework.ai.vectorstore.SimpleVectorStore;
-import org.springframework.ai.vectorstore.VectorStore;
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
-import org.springframework.context.annotation.Bean;
 
 @SpringBootApplication
 public class QuestionAnswering {
@@ -14,9 +10,4 @@ public static void main(String[] args) {
         SpringApplication.run(QuestionAnswering.class, args);
     }
 
-    @Bean
-    VectorStore vectorStore(EmbeddingModel embeddingModel) {
-        return new SimpleVectorStore(embeddingModel);
-    }
-
 }
diff --git a/00-use-cases/question-answering/src/main/resources/application.yml b/00-use-cases/question-answering/src/main/resources/application.yml
@@ -4,6 +4,12 @@ spring:
       chat:
         options:
           model: mistral
+          num-ctx: 4096
       embedding:
         options:
-          model: mistral
+          model: nomic-embed-text
+    vectorstore:
+      pgvector:
+        initialize-schema: true
+        dimensions: 768
+        index-type: hnsw
diff --git a/...es/question-answering/src/test/java/com/thomasvitale/ai/spring/TestQuestionAnswering.java b/...es/question-answering/src/test/java/com/thomasvitale/ai/spring/TestQuestionAnswering.java
@@ -5,6 +5,8 @@
 import org.springframework.boot.test.context.TestConfiguration;
 import org.springframework.boot.testcontainers.service.connection.ServiceConnection;
 import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Profile;
+import org.testcontainers.containers.PostgreSQLContainer;
 import org.testcontainers.ollama.OllamaContainer;
 import org.testcontainers.utility.DockerImageName;
 
@@ -14,6 +16,14 @@ public class TestQuestionAnswering {
     @Bean
     @RestartScope
     @ServiceConnection
+    PostgreSQLContainer<?> pgvectorContainer() {
+        return new PostgreSQLContainer<>(DockerImageName.parse("pgvector/pgvector:pg16"));
+    }
+
+    @Bean
+    @Profile("ollama-image")
+    @RestartScope
+    @ServiceConnection
     OllamaContainer ollama() {
         return new OllamaContainer(DockerImageName.parse("ghcr.io/thomasvitale/ollama-mistral")
                 .asCompatibleSubstituteFor("ollama/ollama"));

diff --git a/00-use-cases/semantic-search/README.md b/00-use-cases/semantic-search/README.md
@@ -1,37 +1,38 @@
 # Semantic Search
 
-Semantic search with LLMs via Ollama.
+Semantic search with LLMs via Ollama and PGVector.
 
 ## Running the application
 
 The application relies on Ollama for providing LLMs. You can either run Ollama locally on your laptop, or rely on the Testcontainers support in Spring Boot to spin up an Ollama service automatically.
+Furthermore, the application relies on the native Testcontainers support in Spring Boot to spin up a PostgreSQL database with the pgvector extension for embeddings.
 
 ### Ollama as a native application
 
 First, make sure you have [Ollama](https://ollama.ai) installed on your laptop.
-Then, use Ollama to run the _mistral_ large language model.
+Then, use Ollama to run the _nomic-embed-text_ large language model.
 
 ```shell
-ollama run mistral
+ollama run nomic-embed-text
 ```
 
 Finally, run the Spring Boot application.
 
 ```shell
-./gradlew bootRun
+./gradlew bootTestRun
 ```
 
 ### Ollama as a dev service with Testcontainers
 
-The application relies on the native Testcontainers support in Spring Boot to spin up an Ollama service with a _mistral_ model at startup time.
+The application relies on the native Testcontainers support in Spring Boot to spin up an Ollama service with a _nomic-embed-text_ model at startup time.
 
 ```shell
-./gradlew bootTestRun
+./gradlew bootTestRun --args='--spring.profiles.active=ollama-image'
 ```
 
 ## Calling the application
 
-You can now call the application that will use Ollama and _mistral_ to perform a semantic search.
+You can now call the application that will use Ollama and _nomic-embed-text_ to perform a semantic search.
 This example uses [httpie](https://httpie.io) to send HTTP requests.
 
 ```shell

diff --git a/00-use-cases/semantic-search/build.gradle b/00-use-cases/semantic-search/build.gradle
@@ -24,6 +24,7 @@ dependencies {
 
     implementation 'org.springframework.boot:spring-boot-starter-web'
     implementation 'org.springframework.ai:spring-ai-ollama-spring-boot-starter'
+    implementation 'org.springframework.ai:spring-ai-pgvector-store-spring-boot-starter'
 
     testAndDevelopmentOnly 'org.springframework.boot:spring-boot-devtools'
 
@@ -33,6 +34,7 @@ dependencies {
     testImplementation 'org.springframework.ai:spring-ai-spring-boot-testcontainers'
     testImplementation 'org.testcontainers:junit-jupiter'
     testImplementation 'org.testcontainers:ollama'
+    testImplementation 'org.testcontainers:postgresql'
 }
 
 tasks.named('test') {

diff --git a/00-use-cases/semantic-search/src/main/java/com/thomasvitale/ai/spring/SemanticSearch.java b/00-use-cases/semantic-search/src/main/java/com/thomasvitale/ai/spring/SemanticSearch.java
@@ -1,11 +1,7 @@
 package com.thomasvitale.ai.spring;
 
-import org.springframework.ai.embedding.EmbeddingModel;
-import org.springframework.ai.vectorstore.SimpleVectorStore;
-import org.springframework.ai.vectorstore.VectorStore;
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
-import org.springframework.context.annotation.Bean;
 
 @SpringBootApplication
 public class SemanticSearch {
@@ -14,9 +10,4 @@ public static void main(String[] args) {
         SpringApplication.run(SemanticSearch.class, args);
     }
 
-    @Bean
-    VectorStore vectorStore(EmbeddingModel embeddingModel) {
-        return new SimpleVectorStore(embeddingModel);
-    }
-
 }
diff --git a/00-use-cases/semantic-search/src/main/resources/application.yml b/00-use-cases/semantic-search/src/main/resources/application.yml
@@ -3,4 +3,9 @@ spring:
     ollama:
       embedding:
         options:
-          model: mistral
+          model: nomic-embed-text
+    vectorstore:
+      pgvector:
+        initialize-schema: true
+        dimensions: 768
+        index-type: hnsw
diff --git a/...svitale/ai/spring/TestClassification.java → ...svitale/ai/spring/TestSemanticSearch.java b/...svitale/ai/spring/TestClassification.java → ...svitale/ai/spring/TestSemanticSearch.java
@@ -5,22 +5,32 @@
 import org.springframework.boot.test.context.TestConfiguration;
 import org.springframework.boot.testcontainers.service.connection.ServiceConnection;
 import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Profile;
+import org.testcontainers.containers.PostgreSQLContainer;
 import org.testcontainers.ollama.OllamaContainer;
 import org.testcontainers.utility.DockerImageName;
 
 @TestConfiguration(proxyBeanMethods = false)
-public class TestClassification {
+public class TestSemanticSearch {
 
     @Bean
     @RestartScope
     @ServiceConnection
+    PostgreSQLContainer<?> pgvectorContainer() {
+        return new PostgreSQLContainer<>(DockerImageName.parse("pgvector/pgvector:pg16"));
+    }
+
+    @Bean
+    @Profile("ollama-image")
+    @RestartScope
+    @ServiceConnection
     OllamaContainer ollama() {
-        return new OllamaContainer(DockerImageName.parse("ghcr.io/thomasvitale/ollama-mistral")
+        return new OllamaContainer(DockerImageName.parse("ghcr.io/thomasvitale/ollama-nomic-embed-text")
                 .asCompatibleSubstituteFor("ollama/ollama"));
     }
 
     public static void main(String[] args) {
-        SpringApplication.from(SemanticSearch::main).with(TestClassification.class).run(args);
+        SpringApplication.from(SemanticSearch::main).with(TestSemanticSearch.class).run(args);
     }
 
 }
diff --git a/...s/structured-data-extraction/src/main/java/com/thomasvitale/ai/spring/PatientJournal.java b/...s/structured-data-extraction/src/main/java/com/thomasvitale/ai/spring/PatientJournal.java
@@ -3,6 +3,13 @@
 import java.util.List;
 
 public record PatientJournal(String fullName, List<Observation> observations, Diagnosis diagnosis) {
-    public record Observation(String type, String content) {}
+    public record Observation(Type type, String content) {}
     public record Diagnosis(String content) {}
+
+    enum Type {
+        BODY_WEIGHT,
+        TEMPERATURE,
+        VITAL_SIGNS,
+        OTHER
+    }
 }
diff --git a/...-extraction/src/main/java/com/thomasvitale/ai/spring/StructuredDataExtractionService.java b/...-extraction/src/main/java/com/thomasvitale/ai/spring/StructuredDataExtractionService.java
@@ -24,7 +24,6 @@ PatientJournal extract(String message) {
                     Extract structured data from the provided text.
                     If you do not know the value of a field asked to extract,
                     do not include any value for the field in the result.
-                    Finally, save the object in the database.
         
                     ---------------------
                     TEXT:

diff --git a/00-use-cases/text-classification/README.md b/00-use-cases/text-classification/README.md
@@ -12,7 +12,7 @@ First, make sure you have [Ollama](https://ollama.ai) installed on your laptop.
 Then, use Ollama to pull the _mistral_ large language model.
 
 ```shell
-ollama pull mistral
+ollama pull mistral-nemo
 ```
 
 Finally, run the Spring Boot application.

diff --git a/...ification/src/test/java/com/thomasvitale/ai/spring/TestTextClassificationApplication.java b/...ification/src/test/java/com/thomasvitale/ai/spring/TestTextClassificationApplication.java
@@ -1,9 +1,22 @@
 package com.thomasvitale.ai.spring;
 
 import org.springframework.boot.SpringApplication;
+import org.springframework.boot.devtools.restart.RestartScope;
+import org.springframework.boot.testcontainers.service.connection.ServiceConnection;
+import org.springframework.context.annotation.Bean;
+import org.testcontainers.ollama.OllamaContainer;
+import org.testcontainers.utility.DockerImageName;
 
 public class TestTextClassificationApplication {
 
+    @Bean
+    @RestartScope
+    @ServiceConnection
+    OllamaContainer ollama() {
+        return new OllamaContainer(DockerImageName.parse("ghcr.io/thomasvitale/ollama-mistral-nemo")
+                .asCompatibleSubstituteFor("ollama/ollama"));
+    }
+
     public static void main(String[] args) {
         SpringApplication.from(TextClassificationApplication::main)
                 .with(TestcontainersConfiguration.class).run(args);

diff --git a/README.md b/README.md
@@ -14,13 +14,13 @@ Samples showing how to build Java applications powered by Generative AI and LLMs
 
 ### 0. Use Cases
 
-| Project                                                                                                                                 | Description                                                    |
-|-----------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------|
-| [chatbot](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/chatbot)                                       | Chatbot using LLMs via Ollama.                                 |
-| [question-answering](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/question-answering)                 | Question answering with documents (RAG) using LLMs via Ollama. |
-| [semantic-search](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/semantic-search)                       | Semantic search using LLMs via Ollama.                         |
-| [structured-data-extraction](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/structured-data-extraction) | Structured data extraction using LLMs via Ollama.              |
-| [text-classification](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/text-classification)               | Text classification using LLMs via Ollama.                     |
+| Project                                                                                                                                 | Description                                                                 |
+|-----------------------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------|
+| [chatbot](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/chatbot)                                       | Chatbot using LLMs via Ollama.                                              |
+| [question-answering](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/question-answering)                 | Question answering with documents (RAG) using LLMs via Ollama and PGVector. |
+| [semantic-search](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/semantic-search)                       | Semantic search using LLMs via Ollama and PGVector.                         |
+| [structured-data-extraction](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/structured-data-extraction) | Structured data extraction using LLMs via Ollama.                           |
+| [text-classification](https://github.com/ThomasVitale/llm-apps-java-spring-ai/tree/main/00-use-cases/text-classification)               | Text classification using LLMs via Ollama.                                  |
 
 ### 1. Chat Completion Models