microsoft--graphrag/packages/graphrag-llm/notebooks/05_caching.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "347b0fc9",
   "metadata": {},
   "source": [
    "# Caching\n",
    "\n",
    "To enabling caching, pass in a `Cache` instance to the `create_completion` or `create_embedding` functions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b0c42f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metrics for: azure/gpt-4o\n",
      "{\n",
      "  \"attempted_request_count\": 3,\n",
      "  \"successful_response_count\": 3,\n",
      "  \"failed_response_count\": 0,\n",
      "  \"failure_rate\": 0.0,\n",
      "  \"requests_with_retries\": 0,\n",
      "  \"retries\": 0,\n",
      "  \"retry_rate\": 0.0,\n",
      "  \"compute_duration_seconds\": 2.8864328861236572,\n",
      "  \"compute_duration_per_response_seconds\": 0.9621442953745524,\n",
      "  \"cached_responses\": 1,\n",
      "  \"streaming_responses\": 0,\n",
      "  \"responses_with_tokens\": 3,\n",
      "  \"prompt_tokens\": 191,\n",
      "  \"completion_tokens\": 59,\n",
      "  \"total_tokens\": 250,\n",
      "  \"tokens_per_response\": 83.33333333333333,\n",
      "  \"responses_with_cost\": 3,\n",
      "  \"input_cost\": 0.0004775,\n",
      "  \"output_cost\": 0.00059,\n",
      "  \"total_cost\": 0.0010675,\n",
      "  \"cost_per_response\": 0.0003558333333333334\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# Copyright (c) 2024 Microsoft Corporation.\n",
    "# Licensed under the MIT License\n",
    "\n",
    "import json\n",
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from graphrag_cache import CacheConfig, CacheType, create_cache\n",
    "from graphrag_llm.completion import LLMCompletion, create_completion\n",
    "from graphrag_llm.config import AuthMethod, ModelConfig\n",
    "from graphrag_storage import StorageConfig, StorageType\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "cache = create_cache()\n",
    "# The above default is equivalent to:\n",
    "cache = create_cache(\n",
    "    CacheConfig(\n",
    "        type=CacheType.Json,\n",
    "        storage=StorageConfig(type=StorageType.File, base_dir=\"cache\"),\n",
    "    )\n",
    ")\n",
    "\n",
    "api_key = os.getenv(\"GRAPHRAG_API_KEY\")\n",
    "model_config = ModelConfig(\n",
    "    model_provider=\"azure\",\n",
    "    model=os.getenv(\"GRAPHRAG_MODEL\", \"gpt-4o\"),\n",
    "    azure_deployment_name=os.getenv(\"GRAPHRAG_MODEL\", \"gpt-4o\"),\n",
    "    api_base=os.getenv(\"GRAPHRAG_API_BASE\"),\n",
    "    api_version=os.getenv(\"GRAPHRAG_API_VERSION\", \"2025-04-01-preview\"),\n",
    "    api_key=api_key,\n",
    "    auth_method=AuthMethod.AzureManagedIdentity if not api_key else AuthMethod.ApiKey,\n",
    ")\n",
    "llm_completion: LLMCompletion = create_completion(model_config, cache=cache)\n",
    "\n",
    "response = llm_completion.completion(\n",
    "    messages=\"What is the capital of France?\",\n",
    ")\n",
    "response = llm_completion.completion(\n",
    "    messages=\"What is the capital of France?\",\n",
    ")\n",
    "\n",
    "print(f\"Metrics for: {llm_completion.metrics_store.id}\")\n",
    "metrics = llm_completion.metrics_store.get_metrics()\n",
    "print(json.dumps(metrics, indent=2))\n",
    "assert metrics[\"cached_responses\"] == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c70a72fc",
   "metadata": {},
   "source": [
    "## Note on the above metrics\n",
    "\n",
    "`cached_responses == 1` since the request was cached by the time the second call was made.\n",
    "\n",
    "The `cached_responses` indicates how many cache hits occurred but the rest of the metrics exist as if a cache was not used. For example, `compute_duration_seconds` and all the token counts and cost counts are as if cache was not used. This is because both the response and metrics are cached and retrieved from the cache when a cache hit occurs. Metrics were designed to give an idea of how long and costly a job would be if there were no cache.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27b026d7",
   "metadata": {},
   "source": [
    "## Tests\n",
    "\n",
    "This is in here because notebooks are being used as integration tests. This ensures objects are being loaded and deserialized from cache properly and the cache is bypassing the rate limiting.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22cc179e",
   "metadata": {},
   "source": [
    "### Test Timing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "efb228ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total time for 100 requests: 0.3867683410644531 seconds\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "from graphrag_llm.config import RateLimitConfig, RateLimitType\n",
    "\n",
    "model_config.rate_limit = RateLimitConfig(\n",
    "    type=RateLimitType.SlidingWindow,\n",
    "    period_in_seconds=60,  # limit requests per minute\n",
    "    requests_per_period=1,  # max 1 request per minute. Without cache this would take forever\n",
    ")\n",
    "llm_completion: LLMCompletion = create_completion(model_config, cache=cache)\n",
    "\n",
    "start_time = time.time()\n",
    "for _ in range(100):\n",
    "    response = llm_completion.completion(\n",
    "        messages=\"What is the capital of France?\",\n",
    "    )\n",
    "end_time = time.time()\n",
    "total_time = end_time - start_time\n",
    "print(f\"Total time for 100 requests: {total_time} seconds\")\n",
    "assert total_time < 5.0  # Ensure that caching is effective"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcf4bf16",
   "metadata": {},
   "source": [
    "### Test Structured Responses\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "21e0e1e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrag_llm.types import LLMCompletionResponse\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "class LocalWeather(BaseModel):\n",
    "    \"\"\"City weather information model.\"\"\"\n",
    "\n",
    "    city: str = Field(description=\"The name of the city\")\n",
    "    temperature: float = Field(description=\"The temperature in Celsius\")\n",
    "    condition: str = Field(description=\"The weather condition description\")\n",
    "\n",
    "\n",
    "class WeatherReports(BaseModel):\n",
    "    \"\"\"Weather information model.\"\"\"\n",
    "\n",
    "    reports: list[LocalWeather] = Field(\n",
    "        description=\"The weather reports for multiple cities\"\n",
    "    )\n",
    "\n",
    "\n",
    "llm_completion.metrics_store.clear_metrics()\n",
    "response: LLMCompletionResponse[WeatherReports] = llm_completion.completion(  # type: ignore\n",
    "    messages=\"It is sunny and 52 degrees fahrenheit in Seattle. It is cloudy and 75 degrees fahrenheit in San Francisco.\",\n",
    "    response_format=WeatherReports,\n",
    ")  # type: ignore\n",
    "response: LLMCompletionResponse[WeatherReports] = llm_completion.completion(  # type: ignore\n",
    "    messages=\"It is sunny and 52 degrees fahrenheit in Seattle. It is cloudy and 75 degrees fahrenheit in San Francisco.\",\n",
    "    response_format=WeatherReports,\n",
    ")  # type: ignore\n",
    "\n",
    "metrics = llm_completion.metrics_store.get_metrics()\n",
    "assert metrics[\"cached_responses\"] == 1, (\n",
    "    f\"Expected 1 cached response, got {metrics['cached_responses']}\"\n",
    ")\n",
    "\n",
    "\n",
    "# Changing the response format should not hit the cache and\n",
    "# instead be a new request and store a new response in the cache.\n",
    "\n",
    "\n",
    "class WeatherReports2(BaseModel):\n",
    "    \"\"\"Weather information model.\"\"\"\n",
    "\n",
    "    local_reports: list[LocalWeather] = Field(\n",
    "        description=\"The weather reports for multiple cities\"\n",
    "    )\n",
    "\n",
    "\n",
    "llm_completion.metrics_store.clear_metrics()\n",
    "# Same request but different response format. Should not hit cache.\n",
    "response: LLMCompletionResponse[WeatherReports2] = llm_completion.completion(\n",
    "    messages=\"It is sunny and 52 degrees fahrenheit in Seattle. It is cloudy and 75 degrees fahrenheit in San Francisco.\",\n",
    "    response_format=WeatherReports2,\n",
    ")  # type: ignore\n",
    "\n",
    "metrics = llm_completion.metrics_store.get_metrics()\n",
    "assert metrics.get(\"cached_responses\", 0) == 0, (\n",
    "    f\"Expected 0 cached responses, got {metrics['cached_responses']}\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}