{ "cells": [ { "cell_type": "markdown", "id": "6498d2b8-d6f9-4bad-9c6f-8c8151675b02", "metadata": {}, "source": [ "![Redis](https://redis.io/wp-content/uploads/2024/04/Logotype.svg?auto=webp&quality=85,75&width=120)\n", "\n", "# RAG with Spring AI and Redis\n", "\n", "This notebook demonstrates how to build a Retrieval-Augmented Generation (RAG) system using Spring AI and Redis. The example focuses on creating a beer recommendation chatbot that can answer questions about beers by retrieving relevant information from a database." ] }, { "cell_type": "markdown", "id": "b0cd181e-fceb-4960-a334-1599bfabbd91", "metadata": {}, "source": [ "## Maven Dependencies\n", "\n", "The notebook requires several dependencies:\n", "\n", "- Spring AI OpenAI: To interact with OpenAI's language models\n", "- Spring AI Transformers: For embedding generation using local models\n", "- Spring AI Redis Store: To use Redis as a vector database\n", "- SLF4J: For logging\n", "- Jedis: Redis client for Java" ] }, { "cell_type": "code", "id": "f0483426-9a2a-4fc1-a184-9ba3343d2bf9", "metadata": { "jupyter": { "is_executing": true } }, "source": [ "%mavenRepo spring_milestones https://repo.spring.io/milestone/ \n", "%maven \"org.springframework.ai:spring-ai-openai:1.0.0-M6\"\n", "%maven \"org.springframework.ai:spring-ai-transformers:1.0.0-M6\"\n", "%maven \"org.springframework.ai:spring-ai-redis-store:1.0.0-M6\"\n", "%maven \"org.slf4j:slf4j-simple:2.0.17\" \n", "%maven \"redis.clients:jedis:5.2.0\"" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "id": "e3b4b75f-dc96-462d-88a3-44b1c469ca2a", "metadata": {}, "source": [ "## Setting up the OpenAI Chat Model\n", "\n", "To run the code below, you need to have your OpenAI API key available in environment variable `OPENAI_API_KEY`." ] }, { "cell_type": "code", "execution_count": 2, "id": "c34b42d5-aa83-48c3-b65b-a858ac60c03d", "metadata": {}, "outputs": [], "source": [ "import org.springframework.ai.openai.OpenAiChatModel;\n", "import org.springframework.ai.openai.OpenAiChatOptions;\n", "import org.springframework.ai.openai.api.OpenAiApi;\n", "\n", "var openAiApi = new OpenAiApi(System.getenv(\"OPENAI_API_KEY\"));\n", "\n", "var openAiChatOptions = OpenAiChatOptions.builder()\n", " .model(\"gpt-3.5-turbo\")\n", " .temperature(0.4)\n", " .maxTokens(200)\n", " .build();\n", "\n", "var chatModel = OpenAiChatModel.builder()\n", " .openAiApi(openAiApi)\n", " .defaultOptions(openAiChatOptions)\n", " .build();" ] }, { "cell_type": "markdown", "id": "70f85ac4-ce9a-4be9-b5bd-23518a0c7e09", "metadata": {}, "source": [ "## Setting up the Embedding Model\n", "\n", "Initializes the transformer-based embedding model. Unlike the chat model which uses OpenAI's API, this embedding model runs locally using the Hugging Face transformer models." ] }, { "cell_type": "code", "execution_count": 3, "id": "0094dc34-3b4b-4b9e-8a10-76bb0a57386f", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[JJava-executor-0] INFO org.springframework.ai.transformers.ResourceCacheService - Create cache root directory: /tmp/spring-ai-onnx-generative\n", "[JJava-executor-0] INFO org.springframework.ai.transformers.ResourceCacheService - Caching the URL [https://raw.githubusercontent.com/spring-projects/spring-ai/main/models/spring-ai-transformers/src/main/resources/onnx/all-MiniLM-L6-v2/tokenizer.json] resource to: /tmp/spring-ai-onnx-generative/4d42ba07-cb22-352f-bb44-beccc8c8c0b7/tokenizer.json\n", "[JJava-executor-0] INFO ai.djl.util.Platform - Found matching platform from: jar:file:/home/jovyan/.ivy2/cache/ai.djl.huggingface/tokenizers/jars/tokenizers-0.30.0.jar!/native/lib/tokenizers.properties\n", "[JJava-executor-0] INFO org.springframework.ai.transformers.ResourceCacheService - Caching the URL [https://github.com/spring-projects/spring-ai/raw/main/models/spring-ai-transformers/src/main/resources/onnx/all-MiniLM-L6-v2/model.onnx] resource to: /tmp/spring-ai-onnx-generative/eb4e1bd7-63c5-301b-8383-5df6a4a2adea/model.onnx\n", "[JJava-executor-0] INFO org.springframework.ai.transformers.TransformersEmbeddingModel - Model input names: input_ids, attention_mask, token_type_ids\n", "[JJava-executor-0] INFO org.springframework.ai.transformers.TransformersEmbeddingModel - Model output names: last_hidden_state\n" ] } ], "source": [ "import org.springframework.ai.transformers.TransformersEmbeddingModel;\n", "\n", "var embeddingModel = new TransformersEmbeddingModel();\n", "embeddingModel.afterPropertiesSet();" ] }, { "cell_type": "markdown", "id": "787c39d1-72ee-429c-8617-3476fc5cc447", "metadata": {}, "source": [ "## Testing the Embedding Model\n", "\n", "Generating vector embeddings for two sample phrases" ] }, { "cell_type": "code", "execution_count": 4, "id": "bc1a02cf-0efc-4480-8d04-bd5d41e50293", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[JJava-executor-0] INFO ai.djl.pytorch.engine.PtEngine - PyTorch graph executor optimizer is enabled, this may impact your inference latency and throughput. See: https://docs.djl.ai/master/docs/development/inference_performance_optimization.html#graph-executor-optimization\n", "[JJava-executor-0] INFO ai.djl.pytorch.engine.PtEngine - Number of inter-op threads is 12\n", "[JJava-executor-0] INFO ai.djl.pytorch.engine.PtEngine - Number of intra-op threads is 12\n" ] } ], "source": [ "List embeddings = embeddingModel.embed(List.of(\"Hello world\", \"World is big\"));" ] }, { "cell_type": "code", "execution_count": 5, "id": "7f42785a-8fd1-415a-8d49-e88c84ceaf21", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "2" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "embeddings.size()" ] }, { "cell_type": "code", "execution_count": 6, "id": "2c0e08b2-cd24-4d47-b752-4a21d1534d23", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[-0.19744644, 0.17766532, 0.03857004, 0.1495222, -0.22542009, -0.918028, 0.38326377, -0.03688945, -0.271742, 0.084521994, 0.40589252, 0.31799775, 0.10991715, -0.15033704, -0.0578956, -0.1542844, 0.1277511, -0.12728858, -0.85726726, -0.100180045, 0.043960992, 0.31126785, 0.018637724, 0.18169005, -0.4846143, -0.16840324, 0.29548055, 0.27559924, -0.01898329, -0.33375576, 0.24035157, 0.12719727, 0.7341182, -0.12793198, -0.06675415, 0.3603812, -0.18827778, -0.52243793, -0.17853652, 0.301802, 0.2693615, -0.48221794, -0.17212732, -0.11880259, 0.054506138, -0.021313868, 0.042054005, 0.22520447, 0.53416646, -0.02169647, -0.30204588, -0.3324908, -0.039310955, 0.030255951, 0.47471577, 0.11088768, 0.03599049, -0.059162557, 0.05172684, -0.21580887, -0.2588888, 0.13753763, -0.03976778, 0.077264294, 0.5730004, -0.41052252, -0.12424426, 0.18107419, -0.29570377, -0.47102028, -0.3762157, -0.0566694, 0.03330949, 0.42123562, -0.19500081, 0.14251879, 0.08297111, 0.15151738, 0.055302583, 0.17305022, 0.30240083, -0.4315744, 0.05667964, 0.170871, 0.10053837, 0.13224423, 0.011074826, 0.00801868, -0.27016994, -0.064108744, -0.65401405, -0.11346026, 0.23059894, 0.012559483, -0.45695782, -0.14536054, 0.5410899, -0.1659703, -0.8304071, 1.3227727, 0.15881175, 0.18389726, 0.17790473, 0.24529731, 0.36788028, 0.1841938, -0.027928434, 0.31898242, -0.21494238, -0.12315938, -0.1623146, -0.16520146, 0.21964264, -0.10004018, 0.3005754, -0.42880356, -0.17901944, 0.12508321, -0.22847626, -0.04917716, 0.15437645, -0.2777267, 0.06568631, 0.16961928, -0.11781378, 0.07504356, 0.16512455, -1.8292688E-32, 0.37099707, -0.103828706, 0.29659325, 0.6985769, 0.16481955, 0.04994966, -0.4038639, -0.09682532, 0.23331007, 0.24119315, 0.14573209, 0.2047131, -0.2814445, 0.012193024, -0.08903271, 0.2905263, -0.2759496, 0.20548306, -0.0232912, 0.5825621, -0.32053158, -0.061168656, 0.064345926, 0.5193481, 0.024250127, 0.20123425, -0.05556667, -0.537552, 0.5317701, 0.045843065, -0.04412724, -0.2982929, -0.07208949, 0.018709056, 0.034438692, 0.043418773, 0.06023024, -0.49448788, -0.40018526, -0.014510898, -0.521009, 0.26851663, 0.29823413, 0.041198455, 0.06244344, -0.029948883, 0.07981756, 0.12580922, 0.19590716, 0.34489778, 6.682277E-4, 0.084367484, -0.40139028, 0.16320959, -0.15807047, 0.061669067, 0.1994718, -0.12878472, 0.05594621, 0.44227248, 0.12363334, 0.65833676, -0.3894322, 0.13607582, -0.091537476, -0.10209247, 0.36878014, 0.18340643, 0.28789037, -0.03386706, -0.1930407, 0.102169015, 0.09491301, 0.36249012, 0.19859105, 0.26614627, 0.5606941, -0.038000442, 0.14435697, -0.44662768, 0.096934825, -0.0054164976, 0.12869316, -0.21907079, 0.548087, -0.030643288, 0.059955206, -0.6599656, -0.075952515, -0.061331585, -0.4759999, 0.41962653, 0.28286183, -0.051509358, -0.548893, 1.927742E-32, 0.7154652, 0.110812716, -0.33345005, -0.20609923, -0.29061896, -0.26150167, -0.47305745, 0.8486894, -0.50637484, 0.34518296, 0.29224205, 0.059004746, 0.80871284, 0.17646644, 0.34952724, -0.30267116, 0.7825679, 0.05262854, -0.09921885, -0.07358193, -0.045787632, -0.29195526, -0.2998041, 0.04348392, -0.08685544, 0.09712923, 0.12181321, 0.11773253, -0.68738264, 0.08282088, 0.15324913, 0.14506459, -0.24484996, 0.038762033, -0.08280242, 0.2592085, -0.5238729, -0.11132506, -0.102130055, -0.3144619, -0.30146742, -0.059897322, -0.29788807, 0.11964548, -0.45797828, -0.06935966, -0.33061957, 0.13273829, -0.045996144, -0.14883682, -0.4578995, -0.11871089, 0.27957174, -0.116765395, -0.28162748, 0.081090145, -0.36435378, -0.044711765, 0.09410101, -0.14707984, 0.07663135, 0.15032242, 0.0571447, 0.36210248, 0.015302703, -0.037698798, 0.09524873, 0.18535785, 0.21729061, -0.20832026, -0.03957802, 9.149015E-4, -0.009355202, -0.15621811, -0.16056955, 0.28451854, -0.1653178, -0.013847964, 0.08461365, 0.05592023, 0.03320237, 0.07723324, 0.031887006, 0.21319377, 0.041419506, 0.22996895, 0.466757, 0.41228518, -0.074770994, -0.24557963, -0.06305952, 0.028048843, -0.052857265, 0.20153615, -0.29226974, -8.999385E-8, -0.5075389, 0.13692492, -0.09299688, 0.18154389, 0.15625265, 0.3004808, -0.26956818, -0.33701032, -0.36198398, 0.23416229, 0.28535756, 0.61020494, -0.42666304, -0.07155929, 0.10520587, 0.22606178, -0.1420139, 0.08313233, -0.21228969, 0.114627264, -2.7827127E-4, 0.056504183, 0.14224814, -0.30042008, 0.16787784, -0.4993352, -0.08303764, 0.14900707, -0.107358016, -0.43641558, 0.20068759, 0.59352744, -0.1606408, 0.07283562, -0.4371048, -0.10681938, 0.14303754, 0.4664252, 0.39377174, -0.36684257, -0.48044774, 0.3514127, -0.19211018, -0.60792434, -0.22953579, 0.18629542, 0.4388187, -0.4181522, 0.0019333661, -0.23406522, -0.43402928, 0.15764633, 0.42736888, 0.10146409, 0.52239466, 0.6312138, 0.0032632276, 0.29472238, -0.083333045, 0.1903145, 0.13625453, -0.13108662, 0.22298925, 0.17298983]" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "float[] e0 = embeddings.get(0);\n", "Arrays.toString(e0);" ] }, { "cell_type": "markdown", "id": "8a85a1da-3ca9-475d-9044-74adce03d7fa", "metadata": {}, "source": [ "## Configuring Redis Vector Store\n", "\n", "Sets up a connection to a Redis server at hostname \"redis-java\" on port 6379\n", "Creates a vector store for storing and retrieving embeddings, with:\n", "\n", "- A Redis index named \"beers\"\n", "- A prefix of \"beer:\" for all keys\n", "- Automatic schema initialization" ] }, { "cell_type": "code", "execution_count": 7, "id": "0e03d272-884f-4fa0-9885-fc3e49466c5a", "metadata": {}, "outputs": [], "source": [ "import redis.clients.jedis.JedisPooled;\n", "import org.springframework.ai.vectorstore.redis.RedisVectorStore;\n", "\n", "var jedisPooled = new JedisPooled(\"redis-java\", 6379);\n", "\n", "var vectorStore = RedisVectorStore.builder(jedisPooled, embeddingModel)\n", " .indexName(\"beers\") \n", " .prefix(\"beer:\") \n", " .initializeSchema(true) \n", " .build();\n", "\n", "vectorStore.afterPropertiesSet();" ] }, { "cell_type": "markdown", "id": "d2f90c67-b58f-4613-be1f-487fd56f3146", "metadata": {}, "source": [ "## Loading Beer Data into Redis\n", "\n", "- Defines the relevant fields to extract from the beer JSON data\n", "- Checks if embeddings are already loaded in Redis by querying the index information\n", "- If not loaded:\n", " - Opens the compressed beer data file\n", " - Creates a JSON reader to parse the file and extract the specified fields\n", " - Adds the documents to the vector store, which automatically:\n", " - Creates embeddings for each document\n", " - Stores both the documents and their embeddings in Redis" ] }, { "cell_type": "code", "execution_count": 8, "id": "1f120966-1e4f-422b-9b84-c8bedb2720fc", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Embeddings already loaded. Skipping\n" ] } ], "source": [ "import java.io.File;\n", "import java.io.FileInputStream;\n", "import java.util.Map;\n", "import java.util.zip.GZIPInputStream;\n", "\n", "import org.springframework.ai.reader.JsonReader;\n", "import org.springframework.core.io.InputStreamResource;\n", "import org.springframework.core.io.FileSystemResource;\n", "\n", "// Define the keys we want to extract from the JSON\n", "String[] KEYS = { \"name\", \"abv\", \"ibu\", \"description\" };\n", "\n", "// Data path\n", "String filePath = \"./resources/beers.json.gz\";\n", "\n", "// Check if embeddings are already loaded\n", "Map indexInfo = vectorStore.getJedis().ftInfo(\"beers\");\n", "long numDocs = (long)indexInfo.getOrDefault(\"num_docs\", \"0\");\n", "if (numDocs > 20000) {\n", " System.out.println(\"Embeddings already loaded. Skipping\");\n", "} else {\n", " System.out.println(\"Creating Embeddings...\");\n", " \n", " // Create a file resource directly from the absolute path\n", " File file = new File(filePath);\n", " \n", " // Create a GZIPInputStream\n", " GZIPInputStream inputStream = new GZIPInputStream(new FileInputStream(file));\n", " InputStreamResource resource = new InputStreamResource(inputStream);\n", " \n", " // Create a JSON reader with fields relevant to our use case\n", " JsonReader loader = new JsonReader(resource, KEYS);\n", " \n", " // Use the VectorStore to insert the documents into Redis\n", " vectorStore.add(loader.get());\n", " \n", " System.out.println(\"Embeddings created.\");\n", "}" ] }, { "cell_type": "markdown", "id": "70a3cd51-b016-4e89-a964-4379ef6de06d", "metadata": {}, "source": [ "## Define the System Prompt\n", "\n", "Here we try to control the behavior of the LLM" ] }, { "cell_type": "code", "execution_count": 9, "id": "480bd7cf-d361-4690-9c75-f17a20ebeffb", "metadata": {}, "outputs": [], "source": [ "String systemPrompt = \"\"\"\n", " You're assisting with questions about products in a beer catalog.\n", " Use the information from the DOCUMENTS section to provide accurate answers.\n", " The answer involves referring to the ABV or IBU of the beer, include the beer name in the response.\n", " If unsure, simply state that you don't know.\n", " \n", " DOCUMENTS:\n", " {documents}\n", " \"\"\";" ] }, { "cell_type": "markdown", "id": "f06b2e70-bf67-49e4-897f-95aaf86f54f0", "metadata": {}, "source": [ "## Setting up the Chat Client with the created ChatModel" ] }, { "cell_type": "code", "execution_count": 10, "id": "df0ae72a-051c-43a6-8354-8a540713b988", "metadata": {}, "outputs": [], "source": [ "import org.springframework.ai.chat.client.ChatClient;\n", "\n", "ChatClient chatClient = ChatClient.builder(chatModel)\n", " .build();" ] }, { "cell_type": "markdown", "id": "346aeb8d-0f1c-4223-95f2-7d5ee0da3bb7", "metadata": {}, "source": [ "## Creating a Query Function\n", "\n", "Encapsulate the RAG logic into a single method" ] }, { "cell_type": "code", "execution_count": 11, "id": "5721b36c-6eab-4967-8d15-f1f547b1999c", "metadata": {}, "outputs": [], "source": [ "import java.util.stream.Collectors;\n", "import org.springframework.ai.chat.model.ChatResponse;\n", "import org.springframework.ai.chat.messages.Message;\n", "import org.springframework.ai.chat.messages.UserMessage;\n", "import org.springframework.ai.chat.prompt.Prompt;\n", "import org.springframework.ai.chat.prompt.SystemPromptTemplate;\n", "import org.springframework.ai.document.Document;\n", "import org.springframework.ai.vectorstore.SearchRequest;\n", "\n", "void ask(String query) {\n", " SearchRequest request = SearchRequest.builder().query(query).topK(10).build();\n", "\n", " // Query Redis for the top K documents most relevant to the input message\n", " List docs = vectorStore.similaritySearch(request);\n", " \n", " String documents = docs.stream() //\n", " .map(Document::getText) //\n", " .collect(Collectors.joining(\"\\n\"));\n", " \n", " SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemPrompt);\n", " Message systemMessage = systemPromptTemplate.createMessage(Map.of(\"documents\", documents));\n", " \n", " UserMessage userMessage = new UserMessage(query);\n", " // Assemble the complete prompt using a template\n", " Prompt prompt = new Prompt(List.of(systemMessage, userMessage));\n", " // Call the chat client with the prompt\n", " ChatResponse chatResponse = chatClient.prompt(prompt).call().chatResponse();\n", " \n", " System.out.println(chatResponse.getResult().getOutput().getText());\n", "}" ] }, { "cell_type": "markdown", "id": "82bcb6e1-e805-47ef-8838-0a62ffaeb0e1", "metadata": {}, "source": [ "## 🍺 Now let's talk about Beers!" ] }, { "cell_type": "code", "execution_count": 12, "id": "997b3010-eb42-41f4-8c19-339a95e4047b", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "A beer that pairs well with smoked meats is the \"Oak Smoker,\" with an ABV of 11.5%. This Smoked Wee Heavy has a wonderfully subtle smoky background and rich malty flavors, making it a perfect pairing for BBQ or enjoying on its own.\n" ] } ], "source": [ "ask(\"What beer pais well with smoked meats?\");" ] }, { "cell_type": "code", "execution_count": 13, "id": "1a3d5322-1eae-43d4-847b-54b40713c4de", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Beer does not typically aid in weight loss as it contains calories. However, lower alcohol content beers like the Airship Cream Ale with an ABV of 4.5 might be a lighter option compared to higher ABV beers.\n" ] } ], "source": [ "ask(\"What beer would make me lose weight?\");" ] }, { "cell_type": "code", "execution_count": null, "id": "082c782c-266a-40f7-a073-e5d1852e6d7a", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Java", "language": "java", "name": "java" }, "language_info": { "codemirror_mode": "java", "file_extension": ".jshell", "mimetype": "text/x-java-source", "name": "Java", "pygments_lexer": "java", "version": "21.0.6+7-Ubuntu-124.04.1" } }, "nbformat": 4, "nbformat_minor": 5 }