spring-ai/spring-ai-openai.ipynb

{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Running LLM inference with Spring AI & Ollama\n",
    "This notebook implement the basic text-to-text generation using [Spring AI](https://spring.io/projects/spring-ai) and OpenAI.\n",
    "You need to an OpenAI API key to run. Rename `openaikey.example.json` to `openaikey.secret.json` and update the OpenAI key\n",
    "\n",
    "Free feel to contribute to add more use cases."
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Install dependencies"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-09T13:35:56.465698Z",
     "start_time": "2025-02-09T13:35:56.391538Z"
    }
   },
   "cell_type": "code",
   "source": [
    "// load version variables\n",
    "%use @file[resources/version.json](currentDir=\".\")"
   ],
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-02-09T13:36:30.407213Z",
     "start_time": "2025-02-09T13:35:56.468679Z"
    }
   },
   "cell_type": "code",
   "source": [
    "USE {\n",
    "    repositories {\n",
    "        maven { url = \"https://repo.spring.io/milestone\" }\n",
    "        mavenCentral()\n",
    "    }\n",
    "    dependencies {\n",
    "        implementation(\"org.springframework.ai:spring-ai-core:$springAiVersion\")\n",
    "        implementation(\"org.springframework.ai:spring-ai-ollama:$springAiVersion\")\n",
    "        implementation(\"org.springframework.ai:spring-ai-openai:$springAiVersion\")\n",
    "\n",
    "        implementation(\"io.projectreactor:reactor-core:$reactorVersion\")\n",
    "        implementation(\"org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:$kotlinCoroutineVersion\")\n",
    "        implementation(\"org.jetbrains.kotlinx:kotlinx-coroutines-reactor:$kotlinCoroutineVersion\")\n",
    "        implementation(\"org.jetbrains.kotlinx:kotlinx-coroutines-reactive:$kotlinCoroutineVersion\")\n",
    "        implementation(\"com.fasterxml.jackson.module:jackson-module-kotlin:$jacksonKotlinModule\")\n",
    "    }\n",
    "    import(\n",
    "        \"kotlinx.coroutines.*\",\n",
    "        \"kotlinx.coroutines.flow.*\",\n",
    "        \"kotlinx.coroutines.reactor.*\",\n",
    "        \"kotlinx.coroutines.reactive.*\",\n",
    "    )\n",
    "}\n",
    "// list the library, if the dependencies doesn't show up, run again and restart the kernel\n",
    "notebook.currentClasspath.flatMap { it.split(\"/|\\\\\\\\\".toRegex()).takeLast(1) }.sorted().joinToString(\"\\n\")"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HdrHistogram-2.2.2.jar\n",
       "LatencyUtils-2.0.3.jar\n",
       "ST4-4.3.4.jar\n",
       "annotations-13.0.jar\n",
       "annotations-23.0.0.jar\n",
       "antlr-runtime-3.5.3.jar\n",
       "antlr4-runtime-4.13.1.jar\n",
       "api-0.12.0-363.jar\n",
       "classmate-1.5.1.jar\n",
       "commons-lang3-3.11.jar\n",
       "context-propagation-1.1.2.jar\n",
       "groovy-4.0.16.jar\n",
       "groovy-json-4.0.16.jar\n",
       "jackson-annotations-2.17.2.jar\n",
       "jackson-annotations-2.18.2.jar\n",
       "jackson-core-2.17.2.jar\n",
       "jackson-core-2.18.2.jar\n",
       "jackson-databind-2.16.1.jar\n",
       "jackson-databind-2.18.2.jar\n",
       "jackson-datatype-jsr310-2.16.1.jar\n",
       "jackson-module-jsonSchema-2.17.2.jar\n",
       "jackson-module-kotlin-2.18.2.jar\n",
       "json-path-5.4.0.jar\n",
       "jsonschema-generator-4.35.0.jar\n",
       "jsonschema-module-jackson-4.35.0.jar\n",
       "jsonschema-module-swagger-2-4.35.0.jar\n",
       "jtokkit-1.1.0.jar\n",
       "jul-to-slf4j-2.0.16.jar\n",
       "kotlin-reflect-1.8.10.jar\n",
       "kotlin-reflect-1.9.23.jar\n",
       "kotlin-script-runtime-1.9.23.jar\n",
       "kotlin-stdlib-1.9.23.jar\n",
       "kotlin-stdlib-2.0.0.jar\n",
       "kotlinx-coroutines-core-jvm-1.9.0.jar\n",
       "kotlinx-coroutines-reactive-1.9.0.jar\n",
       "kotlinx-coroutines-reactor-1.9.0.jar\n",
       "kotlinx-serialization-core-jvm-1.6.3.jar\n",
       "kotlinx-serialization-json-jvm-1.6.3.jar\n",
       "lib-0.12.0-363.jar\n",
       "log4j-api-2.23.1.jar\n",
       "log4j-to-slf4j-2.23.1.jar\n",
       "logback-classic-1.5.12.jar\n",
       "logback-core-1.5.12.jar\n",
       "micrometer-commons-1.13.8.jar\n",
       "micrometer-core-1.13.8.jar\n",
       "micrometer-observation-1.12.12.jar\n",
       "reactive-streams-1.0.4.jar\n",
       "reactor-core-3.7.1.jar\n",
       "rest-assured-common-5.4.0.jar\n",
       "slf4j-api-2.0.12.jar\n",
       "slf4j-api-2.0.3.jar\n",
       "spring-ai-core-1.0.0-M5.jar\n",
       "spring-ai-ollama-1.0.0-M5.jar\n",
       "spring-ai-openai-1.0.0-M5.jar\n",
       "spring-ai-retry-1.0.0-M5.jar\n",
       "spring-aop-6.1.15.jar\n",
       "spring-beans-6.1.15.jar\n",
       "spring-boot-starter-logging-3.3.6.jar\n",
       "spring-context-6.1.15.jar\n",
       "spring-context-support-6.1.15.jar\n",
       "spring-core-6.1.15.jar\n",
       "spring-expression-6.1.15.jar\n",
       "spring-jcl-6.1.15.jar\n",
       "spring-messaging-6.1.15.jar\n",
       "spring-retry-2.0.9.jar\n",
       "spring-web-6.1.15.jar\n",
       "spring-webflux-6.1.15.jar\n",
       "swagger-annotations-2.2.20.jar\n",
       "validation-api-1.1.0.Final.jar"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 12
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Function to measure model performance\n",
    "A sample implementation using stream to measure key metrics\n",
    "- time to first token\n",
    "- input token process rate/s\n",
    "- output token process rate/s"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-09T13:36:30.512809Z",
     "start_time": "2025-02-09T13:36:30.410920Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import org.springframework.ai.chat.model.ChatModel\n",
    "import org.springframework.ai.chat.model.ChatResponse\n",
    "import org.springframework.ai.chat.prompt.Prompt\n",
    "import reactor.core.publisher.Flux\n",
    "import kotlin.time.measureTimedValue\n",
    "\n",
    "data class ModelPerformance(\n",
    "    val timeToFirstTokenInMills: Double,\n",
    "    val totalTimeInMills: Double,\n",
    "    val promptTokens: Long,\n",
    "    val generationTokens: Long,\n",
    "    val inputTokenRatePerSec: Double,\n",
    "    val outputTokenRatePerSec: Double\n",
    ")\n",
    "\n",
    "fun <T> measureModelPerformance(block: suspend () -> T): ModelPerformance {\n",
    "    var ttft: Long?\n",
    "\n",
    "    println(\"=== evaluating model performance ===\")\n",
    "    val timedResp = measureTimedValue {\n",
    "        runBlocking {\n",
    "            val startTime = System.nanoTime()\n",
    "            ttft = System.nanoTime() - startTime\n",
    "            val runnable = block()\n",
    "            if (runnable is Flow<*>) {\n",
    "                runnable\n",
    "                    .onStart {\n",
    "                        ttft = System.nanoTime() - startTime\n",
    "                    }\n",
    "                    .last()\n",
    "            } else {\n",
    "                runnable!!\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    println(\"\\n=== model performance ===\")\n",
    "    val resp = timedResp.value\n",
    "\n",
    "    if (resp is ChatResponse) {\n",
    "        val totalTime = timedResp.duration.inWholeMilliseconds\n",
    "        val timeToFirstTokenInMills = ttft!! / 1_000_000.0\n",
    "        return ModelPerformance(\n",
    "            timeToFirstTokenInMills = timeToFirstTokenInMills,\n",
    "            totalTimeInMills = timedResp.duration.inWholeMilliseconds.toDouble(),\n",
    "            promptTokens = resp.metadata.usage.promptTokens,\n",
    "            generationTokens = resp.metadata.usage.generationTokens,\n",
    "            inputTokenRatePerSec = resp.metadata.usage.promptTokens * 1000.0 / timeToFirstTokenInMills,\n",
    "            outputTokenRatePerSec = resp.metadata.usage.generationTokens * 1000.0 / (totalTime - timeToFirstTokenInMills)\n",
    "        )\n",
    "    }\n",
    "    throw IllegalArgumentException(\"not support type\")\n",
    "}"
   ],
   "outputs": [],
   "execution_count": 13
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Create GPT4o-mini chat model\n",
    "\n",
    "### Load OpenAI Key into Kotlin Notebook\n",
    "Rename `openaikey.example.json` to `openaikey.secret.json` and update the OpenAI key\n"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-09T13:36:30.542951Z",
     "start_time": "2025-02-09T13:36:30.515721Z"
    }
   },
   "cell_type": "code",
   "source": [
    "// Load openaikey.json into `openAiKey`\n",
    "%use @file[resources/openaikey.secret.json](currentDir=\".\")"
   ],
   "outputs": [],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-09T13:36:32.609776Z",
     "start_time": "2025-02-09T13:36:30.545685Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import org.springframework.ai.chat.prompt.Prompt\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 model = OpenAiChatModel(\n",
    "    OpenAiApi(openAiKey),\n",
    "    OpenAiChatOptions.builder()\n",
    "        .streamUsage(true)\n",
    "        .model(OpenAiApi.ChatModel.GPT_4_O_MINI)\n",
    "        .temperature(0.7)\n",
    "        .build()\n",
    ")\n",
    "\n",
    "measureModelPerformance {\n",
    "    model.call(Prompt(\"tell me 5 jokes\"))?.also {\n",
    "        print(it.result?.output?.content)\n",
    "    }\n",
    "}"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== evaluating model performance ===\n",
      "Sure! Here are five jokes for you:\n",
      "\n",
      "1. Why don't scientists trust atoms?\n",
      "   Because they make up everything!\n",
      "\n",
      "2. What did the ocean say to the beach?\n",
      "   Nothing, it just waved!\n",
      "\n",
      "3. Why did the scarecrow win an award?\n",
      "   Because he was outstanding in his field!\n",
      "\n",
      "4. How does a penguin build its house?\n",
      "   Igloos it together!\n",
      "\n",
      "5. Why did the bicycle fall over?\n",
      "   Because it was two-tired!\n",
      "\n",
      "Hope these made you smile!\n",
      "=== model performance ===\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ModelPerformance(timeToFirstTokenInMills=8.3E-5, totalTimeInMills=2014.0, promptTokens=12, generationTokens=103, inputTokenRatePerSec=1.4457831325301206E8, outputTokenRatePerSec=51.14200806593181)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 15
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Using"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-09T13:36:38.270631Z",
     "start_time": "2025-02-09T13:36:32.618661Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import org.springframework.ai.chat.client.ChatClient\n",
    "\n",
    "val chatModel = ChatClient.builder(model)\n",
    "    .defaultSystem(\"You are Peter, a helpful AI assistant.\")\n",
    "    .build()\n",
    "\n",
    "measureModelPerformance {\n",
    "    chatModel.prompt()\n",
    "        .user(\"Tell me 5 jokes.\")\n",
    "        .call()\n",
    "        .chatResponse()\n",
    "        .also {\n",
    "            print(it?.result?.output?.content)\n",
    "        }\n",
    "}"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== evaluating model performance ===\n",
      "Sure! Here are five jokes for you:\n",
      "\n",
      "1. Why don’t skeletons fight each other?  \n",
      "   Because they don’t have the guts!\n",
      "\n",
      "2. What do you call fake spaghetti?  \n",
      "   An impasta!\n",
      "\n",
      "3. Why did the scarecrow win an award?  \n",
      "   Because he was outstanding in his field!\n",
      "\n",
      "4. How does a penguin build its house?  \n",
      "   Igloos it together!\n",
      "\n",
      "5. Why was the math book sad?  \n",
      "   Because it had too many problems!\n",
      "\n",
      "Hope these made you smile!\n",
      "=== model performance ===\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ModelPerformance(timeToFirstTokenInMills=4.2E-5, totalTimeInMills=5607.0, promptTokens=26, generationTokens=109, inputTokenRatePerSec=6.190476190476191E8, outputTokenRatePerSec=19.439985877738433)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-09T13:36:40.255417Z",
     "start_time": "2025-02-09T13:36:38.279616Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import com.fasterxml.jackson.databind.ObjectMapper\n",
    "import com.fasterxml.jackson.module.kotlin.jsonMapper\n",
    "import com.fasterxml.jackson.module.kotlin.kotlinModule\n",
    "import org.springframework.ai.chat.client.ChatClient\n",
    "import org.springframework.ai.model.function.FunctionCallback\n",
    "import org.springframework.ai.model.function.DefaultFunctionCallbackBuilder\n",
    "import reactor.core.scheduler.Schedulers\n",
    "import java.util.function.Function;\n",
    "\n",
    "data class GetWeatherRequest(\n",
    "    val date: String,\n",
    "    val city: String,\n",
    "    val country: String\n",
    ")\n",
    "\n",
    "data class GetWeatherResponse(\n",
    "    val date: String,\n",
    "    val city: String,\n",
    "    val country: String,\n",
    "    val unit: String,\n",
    "    val tempature: String\n",
    ")\n",
    "\n",
    "val getWeatherAPI =\n",
    "    DefaultFunctionCallbackBuilder().function(\"GetWeatherTool\", Function<GetWeatherRequest, GetWeatherResponse> { req ->\n",
    "        println(\"=== calling weather API with date=${req.date}, city=${req.city}, country=${req.country} ===\")\n",
    "        // mock response\n",
    "        GetWeatherResponse(\n",
    "            date = req.date,\n",
    "            city = req.city,\n",
    "            country = req.country,\n",
    "            unit = \"celsius\",\n",
    "            tempature = \"24.0\"\n",
    "        )\n",
    "    })\n",
    "        .description(\"Get weather tool with city, country, and date\")\n",
    "        .inputType(GetWeatherRequest::class.java)\n",
    "        .schemaType(FunctionCallback.SchemaType.JSON_SCHEMA)\n",
    "        .objectMapper(jsonMapper {\n",
    "            addModule(kotlinModule())\n",
    "        })\n",
    "        .build()\n",
    "\n",
    "\n",
    "val chatModel = ChatClient.builder(model)\n",
    "    .defaultSystem(\"You are Peter, a helpful AI assistant.\")\n",
    "    .defaultFunctions(getWeatherAPI)\n",
    "    .build()\n",
    "\n",
    "measureModelPerformance {\n",
    "    chatModel.prompt(\"What's today wether in Seattle on 2025-01-01?\")\n",
    "        .call()\n",
    "        .chatResponse()\n",
    "        .also {\n",
    "            print(it?.result?.output?.content)\n",
    "        }\n",
    "}"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== evaluating model performance ===\n",
      "=== calling weather API with date=2025-01-01, city=Seattle, country=US ===\n",
      "The weather in Seattle on January 1, 2025, is expected to be 24.0°C.\n",
      "=== model performance ===\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ModelPerformance(timeToFirstTokenInMills=8.3E-5, totalTimeInMills=1782.0, promptTokens=229, generationTokens=54, inputTokenRatePerSec=2.7590361445783134E9, outputTokenRatePerSec=30.30303171445097)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 17
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-09T13:36:40.259601Z",
     "start_time": "2025-02-09T13:36:40.257952Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "outputs": [],
   "execution_count": null
  }
 ],
 "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": ""
  },
  "ktnbPluginMetadata": {
   "projectLibraries": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}