notebooks/spring-ai/tutorials/3. Streaming.ipynb

{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Streaming\n",
    "\n",
    "Streaming is a capability that allows AI providers to deliver responses as tokens are generated in real-time,\n",
    "rather than waiting for the complete response to be finished.\n",
    "You can visualize it as filling a bucket with water — streaming lets you watch the water level rising in real-time,\n",
    "instead of placing the bucket somewhere, walking away, and only returning once it's completely full.\n",
    "\n",
    "Streaming significantly improves the user experience,\n",
    "as people can immediately see the response forming rather than staring at a loading indicator.\n",
    "This is especially valuable since synchronous requests can take quite a long time to generate.\n",
    "\n",
    "Let's implement streaming responses with Kotlin and Spring AI.\n",
    "For this, we'll need to add dependencies.\n",
    "In addition to `spring-ai-openai`, we'll need the Kotlin coroutines library to work with `Flow`."
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-24T14:24:03.058187Z",
     "start_time": "2025-03-24T14:24:01.930519Z"
    }
   },
   "cell_type": "code",
   "source": [
    "%useLatestDescriptors\n",
    "%use coroutines\n",
    "@file:DependsOn(\"org.jetbrains.kotlinx:kotlinx-coroutines-reactive:1.10.1\")\n",
    "\n",
    "@file:DependsOn(\"org.springframework.ai:spring-ai-openai-spring-boot-starter:1.0.0-M6\")"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "As in the previous notebooks, we'll need an API key.\n",
    "\n",
    "To use the model, we need to provide an API key.\n",
    "\n",
    "You can obtain this API key from\n",
    "[console.anthropic.com](https://console.anthropic.com/settings/keys)\n",
    "for Anthropic models or from\n",
    "[platform.openai.com](https://platform.openai.com/api-keys)\n",
    "for OpenAI models.\n",
    "\n",
    "Then add the generated API key to your environment variables:\n",
    "\n",
    "[MacOS/Linux]\n",
    "```bash\n",
    "export ANTHROPIC_API_KEY=<INSERT KEY HERE> # for Anthropic\n",
    "export OPENAI_API_KEY=<INSERT KEY HERE> # for OpenAI\n",
    "\n",
    "```\n",
    "\n",
    "[Windows]\n",
    "```shell\n",
    "set ANTHROPIC_API_KEY=<INSERT KEY HERE> # for Anthropic\n",
    "set OPENAI_API_KEY=<INSERT KEY HERE> # for OpenAI\n",
    "```\n",
    "\n",
    "Let's retrieve the API key from environment variables:"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-24T14:24:35.051135Z",
     "start_time": "2025-03-24T14:24:34.972461Z"
    }
   },
   "cell_type": "code",
   "source": "val apiKey = System.getenv(\"OPENAI_API_KEY\") ?: \"YOUR_OPENAI_API_KEY\"",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Now let's create a `ChatModel` and pass in the `ChatOptions`:"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-24T14:24:37.619240Z",
     "start_time": "2025-03-24T14:24:37.313249Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import org.springframework.ai.chat.client.ChatClient\n",
    "import org.springframework.ai.openai.OpenAiChatModel\n",
    "import org.springframework.ai.openai.OpenAiChatOptions\n",
    "import org.springframework.ai.openai.api.OpenAiApi\n",
    "\n",
    "val openAiApi = OpenAiApi.builder().apiKey(apiKey).build()\n",
    "val openAiOptions = OpenAiChatOptions.builder()\n",
    "    .model(OpenAiApi.ChatModel.GPT_4_O_MINI)\n",
    "    .temperature(0.7)\n",
    "    .build()\n",
    "\n",
    "val chatCompletion = OpenAiChatModel.builder()\n",
    "    .openAiApi(openAiApi)\n",
    "    .defaultOptions(openAiOptions)\n",
    "    .build()"
   ],
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Now let's ask the model to generate a short fairy tale about rabbits:"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-24T14:25:09.017681Z",
     "start_time": "2025-03-24T14:25:00.102032Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import kotlinx.coroutines.reactive.asFlow\n",
    "\n",
    "val response: Flow<String> = chatCompletion.stream(\"Generate me 10 sentences of a fairy tale about rabbits\").asFlow()\n",
    "\n",
    "runBlocking {\n",
    "    response.collect { print(it) }\n",
    "}"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Once upon a time, in a lush green meadow filled with wildflowers, there lived a clever rabbit named Ruby. Ruby had big dreams of discovering the legendary Golden Carrot, said to grant wisdom to any rabbit who found it. One sunny morning, she gathered her friends—a shy bunny named Benny and a brave hare named Hazel—to embark on an adventure. They set off through the Whispering Woods, where the trees shared secrets of the past with the gentle breeze.\n",
      "\n",
      "As they ventured deeper, they encountered a wise old tortoise named Tobias, who warned them of the tricky fox that guarded the Golden Carrot. Undeterred, Ruby devised a clever plan to outsmart the fox, using their unique talents to create a distraction. When the fox came to investigate, Benny hopped in one direction, while Hazel darted in the other, leaving the fox bewildered and chasing shadows. \n",
      "\n",
      "Finally, they reached the shimmering field where the Golden Carrot was said to grow, surrounded by sparkling fireflies that danced in the twilight. With hearts pounding, they approached the radiant carrot, and Ruby gently plucked it from the earth, feeling its warm glow. As the sun set, the friends shared their dreams and wishes, realizing that the true magic lay not in the carrot itself, but in the bond they had formed during their adventure. And so, they returned home, wiser and closer than ever, knowing that friendship was the greatest treasure of all."
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "As you can see, the response starts printing immediately.\n",
    "If this weren't a streaming response,\n",
    "we would be waiting approximately 8 seconds (judging by the cell execution time) with nothing happening.\n",
    "\n",
    "Now let's do the same thing using the `ChatClient`:"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-03-24T14:25:57.222546Z",
     "start_time": "2025-03-24T14:25:51.921755Z"
    }
   },
   "cell_type": "code",
   "source": [
    "val chatClient = ChatClient.create(chatCompletion)\n",
    "\n",
    "val response: Flow<String> = chatClient\n",
    "    .prompt(\"Generate me 10 sentences of a fairy tale about rabbits\")\n",
    "    .stream()\n",
    "    .content()\n",
    "    .asFlow()\n",
    "\n",
    "runBlocking {\n",
    "    response.collect { print(it) }\n",
    "}"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Once upon a time, in a lush green meadow where wildflowers danced in the breeze, there lived a clever little rabbit named Oliver. Oliver had the softest fur and the brightest eyes, and he dreamed of finding the legendary Golden Carrot said to grant wishes to those pure of heart. One sunny morning, he gathered his friends—a brave rabbit named Bella and a wise old tortoise named Timothy—to embark on an adventure to find this magical treasure.\n",
      "\n",
      "As they journeyed through the enchanted forest, they encountered a mischievous fox who tried to trick them into revealing their quest. But Oliver, quick on his feet, cleverly spun a tale about a hidden treasure guarded by a fierce dragon, which sent the fox scampering away in fear. The trio continued on their path, facing challenges like crossing a bubbling brook and climbing a steep hill, but they encouraged each other with laughter and songs.\n",
      "\n",
      "After what felt like days of travel, they finally stumbled upon a sparkling glade where the Golden Carrot grew atop a shimmering pedestal. But as they approached it, they were met by a gentle spirit who warned them that only those who demonstrated true friendship and bravery could claim the carrot. Together, Oliver, Bella, and Timothy shared stories of their adventures and the bond they had formed, proving their worthiness.\n",
      "\n",
      "With a warm smile, the spirit granted them each a wish, and as they returned home, they discovered that their true treasure was not the Golden Carrot but the friendship they had forged along the way. From that day on, they lived happily in their meadow, sharing their stories and helping other animals in need, spreading kindness and joy wherever they went. And so, the tale of Oliver and his friends became a cherished legend, reminding all who heard it that the greatest magic lies in the bonds of friendship."
     ]
    }
   ],
   "execution_count": 5
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "name": "kotlin",
   "version": "1.9.23",
   "mimetype": "text/x-kotlin",
   "file_extension": ".kt",
   "pygments_lexer": "kotlin",
   "codemirror_mode": "text/x-kotlin",
   "nbconvert_exporter": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}