feat: Add HybridCache, remove MemoryCache (#73)

getcellm · Dec 3, 2024 · e2c5c05 · e2c5c05
1 parent f1f3735
commit e2c5c05
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 47 deletions.
diff --git a/src/Cellm/Cellm.csproj b/src/Cellm/Cellm.csproj
@@ -26,7 +26,6 @@
     <PackageReference Include="Microsoft.Extensions.AI.Ollama" Version="9.0.1-preview.1.24570.5" />
     <PackageReference Include="Microsoft.Extensions.AI.OpenAI" Version="9.0.1-preview.1.24570.5" />
     <PackageReference Include="Microsoft.Extensions.Caching.Hybrid" Version="9.0.0-preview.9.24556.5" />
-    <PackageReference Include="Microsoft.Extensions.Caching.Memory" Version="9.0.0" />
     <PackageReference Include="Microsoft.Extensions.Configuration" Version="9.0.0" />
     <PackageReference Include="Microsoft.Extensions.Configuration.Json" Version="9.0.0" />
     <PackageReference Include="Microsoft.Extensions.DependencyInjection" Version="9.0.0" />

diff --git a/src/Cellm/Models/ModelRequestBehavior/CachingBehavior.cs b/src/Cellm/Models/ModelRequestBehavior/CachingBehavior.cs
@@ -1,51 +1,32 @@
-using System.Security.Cryptography;
-using System.Text;
-using System.Text.Json;
+using System.Text.Json;
 using Cellm.Services.Configuration;
 using MediatR;
-using Microsoft.Extensions.Caching.Memory;
+using Microsoft.Extensions.Caching.Hybrid;
 using Microsoft.Extensions.Options;
 
 namespace Cellm.Models.ModelRequestBehavior;
 
-internal class CachingBehavior<TRequest, TResponse> : IPipelineBehavior<TRequest, TResponse>
+internal class CachingBehavior<TRequest, TResponse>(HybridCache cache, IOptions<CellmConfiguration> cellmConfiguration) : IPipelineBehavior<TRequest, TResponse>
     where TRequest : IModelRequest<TResponse>
     where TResponse : IModelResponse
 {
-    private readonly IMemoryCache _memoryCache;
-    private readonly MemoryCacheEntryOptions _memoryCacheEntryOptions;
-
-    public CachingBehavior(IMemoryCache memoryCache, IOptions<CellmConfiguration> _cellmConfiguration)
+    private readonly HybridCacheEntryOptions _cacheEntryOptions = new()
     {
-        _memoryCache = memoryCache;
-        _memoryCacheEntryOptions = new()
-        {
-            SlidingExpiration = TimeSpan.FromSeconds(_cellmConfiguration.Value.CacheTimeoutInSeconds)
-        };
-    }
+        Expiration = TimeSpan.FromSeconds(cellmConfiguration.Value.CacheTimeoutInSeconds)
+    };
 
     public async Task<TResponse> Handle(TRequest request, RequestHandlerDelegate<TResponse> next, CancellationToken cancellationToken)
     {
-        var key = GetKey(request);
-
-        if (_memoryCache.TryGetValue(key, out object? value) && value is TResponse response)
+        if (cellmConfiguration.Value.EnableCache)
         {
-            return response;
+            return await cache.GetOrCreateAsync(
+                JsonSerializer.Serialize(request.Prompt),
+                async cancel => await next(),
+                options: _cacheEntryOptions,
+                cancellationToken: cancellationToken
+            );
         }
 
-        response = await next();
-
-        _memoryCache.Set(key, response, _memoryCacheEntryOptions);
-
-        return response;
-    }
-
-    private static string GetKey<T>(T key)
-    {
-        var json = JsonSerializer.Serialize(key);
-        var bytes = SHA256.HashData(Encoding.UTF8.GetBytes(json));
-        var hash = Convert.ToHexString(bytes);
-
-        return hash;
+        return await next();
     }
 }
diff --git a/src/Cellm/Services/Configuration/CellmConfiguration.cs b/src/Cellm/Services/Configuration/CellmConfiguration.cs
@@ -14,6 +14,8 @@ public class CellmConfiguration
 
     public int CacheTimeoutInSeconds { get; init; }
 
+    public bool EnableCache { get; init; }
+
     public bool EnableTools { get; init; }
 }
 
diff --git a/src/Cellm/appsettings.json b/src/Cellm/appsettings.json
@@ -31,6 +31,7 @@
         "DefaultTemperature": 0,
         "HttpTimeoutInSeconds": 600,
         "CacheTimeoutInSeconds": 3600,
+        "EnableCache": true,
         "EnableTools": true
     },
     "RateLimiterConfiguration": {

diff --git a/src/Cellm/packages.lock.json b/src/Cellm/packages.lock.json
@@ -83,19 +83,6 @@
           "Microsoft.Extensions.Options": "9.0.0"
         }
       },
-      "Microsoft.Extensions.Caching.Memory": {
-        "type": "Direct",
-        "requested": "[9.0.0, )",
-        "resolved": "9.0.0",
-        "contentHash": "zbnPX/JQ0pETRSUG9fNPBvpIq42Aufvs15gGYyNIMhCun9yhmWihz0WgsI7bSDPjxWTKBf8oX/zv6v2uZ3W9OQ==",
-        "dependencies": {
-          "Microsoft.Extensions.Caching.Abstractions": "9.0.0",
-          "Microsoft.Extensions.DependencyInjection.Abstractions": "9.0.0",
-          "Microsoft.Extensions.Logging.Abstractions": "9.0.0",
-          "Microsoft.Extensions.Options": "9.0.0",
-          "Microsoft.Extensions.Primitives": "9.0.0"
-        }
-      },
       "Microsoft.Extensions.Configuration": {
         "type": "Direct",
         "requested": "[9.0.0, )",
@@ -298,6 +285,18 @@
           "Microsoft.Extensions.Primitives": "9.0.0"
         }
       },
+      "Microsoft.Extensions.Caching.Memory": {
+        "type": "Transitive",
+        "resolved": "9.0.0",
+        "contentHash": "zbnPX/JQ0pETRSUG9fNPBvpIq42Aufvs15gGYyNIMhCun9yhmWihz0WgsI7bSDPjxWTKBf8oX/zv6v2uZ3W9OQ==",
+        "dependencies": {
+          "Microsoft.Extensions.Caching.Abstractions": "9.0.0",
+          "Microsoft.Extensions.DependencyInjection.Abstractions": "9.0.0",
+          "Microsoft.Extensions.Logging.Abstractions": "9.0.0",
+          "Microsoft.Extensions.Options": "9.0.0",
+          "Microsoft.Extensions.Primitives": "9.0.0"
+        }
+      },
       "Microsoft.Extensions.Compliance.Abstractions": {
         "type": "Transitive",
         "resolved": "9.0.0",