Skip to content

Commit 1bf7b42

Browse files
authored
feat: add embedding generators and cache (#1111)
## Summary Implements embedding generation infrastructure with caching support: ### Core Components - **IEmbeddingGenerator**: Core interface for all embedding providers - **EmbeddingConstants**: Known model dimensions (OpenAI, Ollama, HuggingFace models) - **CachedEmbeddingGenerator**: Decorator pattern for transparent caching ### Cache System - **IEmbeddingCache**: Cache interface with SQLite implementation - **EmbeddingCacheKey**: SHA256 hashing (privacy: input text never stored) - **CachedEmbedding**: Model with provider/model metadata and timestamps - **CacheModes**: ReadWrite, ReadOnly, WriteOnly options - **Composite primary key**: provider, model, dimensions, normalized, text_hash for easier debugging and selective queries ### Embedding Providers - **OllamaEmbeddingGenerator**: Local models via Ollama API - **OpenAIEmbeddingGenerator**: OpenAI API (text-embedding-3-small/large, ada-002) - **AzureOpenAIEmbeddingGenerator**: Azure OpenAI deployments - **HuggingFaceEmbeddingGenerator**: HuggingFace Inference API ### Security - Input text is **never stored** in cache - only SHA256 hash - Cache key includes: hash + provider + model + dimensions + normalization flag ## Test Plan - [x] 133 new tests covering all components - [x] Test coverage: 86.33% (above 80% threshold) - [x] All 770 tests pass (0 skipped) - [x] `format.sh` passes - [x] `build.sh` passes with 0 warnings, 0 errors - [x] `coverage.sh` passes
1 parent 0c05ff7 commit 1bf7b42

28 files changed

+4075
-2
lines changed

docs

Submodule docs updated from 03362c9 to 23845bf

src/Core/Config/Embeddings/EmbeddingsConfig.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ namespace KernelMemory.Core.Config.Embeddings;
1212
[JsonDerivedType(typeof(OllamaEmbeddingsConfig), typeDiscriminator: "ollama")]
1313
[JsonDerivedType(typeof(OpenAIEmbeddingsConfig), typeDiscriminator: "openai")]
1414
[JsonDerivedType(typeof(AzureOpenAIEmbeddingsConfig), typeDiscriminator: "azureOpenAI")]
15+
[JsonDerivedType(typeof(HuggingFaceEmbeddingsConfig), typeDiscriminator: "huggingFace")]
1516
public abstract class EmbeddingsConfig : IValidatable
1617
{
1718
/// <summary>
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
using System.Text.Json.Serialization;
3+
using KernelMemory.Core.Config.Enums;
4+
using KernelMemory.Core.Config.Validation;
5+
using KernelMemory.Core.Embeddings;
6+
7+
namespace KernelMemory.Core.Config.Embeddings;
8+
9+
/// <summary>
10+
/// HuggingFace Inference API embeddings provider configuration.
11+
/// Supports the serverless Inference API for embedding models.
12+
/// </summary>
13+
public sealed class HuggingFaceEmbeddingsConfig : EmbeddingsConfig
14+
{
15+
/// <inheritdoc />
16+
[JsonIgnore]
17+
public override EmbeddingsTypes Type => EmbeddingsTypes.HuggingFace;
18+
19+
/// <summary>
20+
/// HuggingFace model name (e.g., "sentence-transformers/all-MiniLM-L6-v2", "BAAI/bge-base-en-v1.5").
21+
/// </summary>
22+
[JsonPropertyName("model")]
23+
public string Model { get; set; } = EmbeddingConstants.DefaultHuggingFaceModel;
24+
25+
/// <summary>
26+
/// HuggingFace API key (token).
27+
/// Can also be set via HF_TOKEN environment variable.
28+
/// </summary>
29+
[JsonPropertyName("apiKey")]
30+
public string? ApiKey { get; set; }
31+
32+
/// <summary>
33+
/// HuggingFace Inference API base URL.
34+
/// Default: https://api-inference.huggingface.co
35+
/// Can be changed for custom inference endpoints.
36+
/// </summary>
37+
[JsonPropertyName("baseUrl")]
38+
public string BaseUrl { get; set; } = EmbeddingConstants.DefaultHuggingFaceBaseUrl;
39+
40+
/// <inheritdoc />
41+
public override void Validate(string path)
42+
{
43+
if (string.IsNullOrWhiteSpace(this.Model))
44+
{
45+
throw new ConfigException($"{path}.Model", "HuggingFace model name is required");
46+
}
47+
48+
if (string.IsNullOrWhiteSpace(this.ApiKey))
49+
{
50+
throw new ConfigException($"{path}.ApiKey", "HuggingFace API key is required");
51+
}
52+
53+
if (string.IsNullOrWhiteSpace(this.BaseUrl))
54+
{
55+
throw new ConfigException($"{path}.BaseUrl", "HuggingFace base URL is required");
56+
}
57+
58+
if (!Uri.TryCreate(this.BaseUrl, UriKind.Absolute, out _))
59+
{
60+
throw new ConfigException($"{path}.BaseUrl",
61+
$"Invalid HuggingFace base URL: {this.BaseUrl}");
62+
}
63+
}
64+
}
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
using System.Text.Json.Serialization;
3+
4+
namespace KernelMemory.Core.Config.Enums;
5+
6+
/// <summary>
7+
/// Modes for embedding cache operations.
8+
/// Controls whether the cache reads, writes, or both.
9+
/// </summary>
10+
[JsonConverter(typeof(JsonStringEnumConverter))]
11+
public enum CacheModes
12+
{
13+
/// <summary>
14+
/// Both read from and write to cache (default).
15+
/// Cache hits return stored embeddings, misses are generated and stored.
16+
/// </summary>
17+
ReadWrite,
18+
19+
/// <summary>
20+
/// Only read from cache, never write.
21+
/// Useful for read-only deployments or when cache is pre-populated.
22+
/// </summary>
23+
ReadOnly,
24+
25+
/// <summary>
26+
/// Only write to cache, never read.
27+
/// Useful for warming up a cache without affecting current behavior.
28+
/// </summary>
29+
WriteOnly
30+
}

src/Core/Config/Enums/EmbeddingsTypes.cs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,5 +16,8 @@ public enum EmbeddingsTypes
1616
OpenAI,
1717

1818
/// <summary>Azure OpenAI Service</summary>
19-
AzureOpenAI
19+
AzureOpenAI,
20+
21+
/// <summary>Hugging Face Inference API</summary>
22+
HuggingFace
2023
}
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
using System.Diagnostics.CodeAnalysis;
3+
4+
namespace KernelMemory.Core.Embeddings.Cache;
5+
6+
/// <summary>
7+
/// Represents a cached embedding vector.
8+
/// </summary>
9+
public sealed class CachedEmbedding
10+
{
11+
/// <summary>
12+
/// The embedding vector.
13+
/// </summary>
14+
[SuppressMessage("Performance", "CA1819:Properties should not return arrays",
15+
Justification = "Embedding vectors are read-only after creation and passed to storage layer")]
16+
public required float[] Vector { get; init; }
17+
}
Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
using System.Security.Cryptography;
3+
using System.Text;
4+
5+
namespace KernelMemory.Core.Embeddings.Cache;
6+
7+
/// <summary>
8+
/// Cache key for embeddings. Uniquely identifies an embedding by provider, model,
9+
/// dimensions, normalization state, and content hash.
10+
/// The input text is NOT stored - only a SHA256 hash is used for security.
11+
/// </summary>
12+
public sealed class EmbeddingCacheKey
13+
{
14+
/// <summary>
15+
/// Provider type name (e.g., "OpenAI", "Ollama", "AzureOpenAI", "HuggingFace").
16+
/// </summary>
17+
public required string Provider { get; init; }
18+
19+
/// <summary>
20+
/// Model name (e.g., "text-embedding-ada-002", "qwen3-embedding").
21+
/// </summary>
22+
public required string Model { get; init; }
23+
24+
/// <summary>
25+
/// Vector dimensions produced by this model.
26+
/// </summary>
27+
public required int VectorDimensions { get; init; }
28+
29+
/// <summary>
30+
/// Whether the vectors are normalized.
31+
/// </summary>
32+
public required bool IsNormalized { get; init; }
33+
34+
/// <summary>
35+
/// Length of the original text in characters.
36+
/// Used as an additional collision prevention measure.
37+
/// </summary>
38+
public required int TextLength { get; init; }
39+
40+
/// <summary>
41+
/// SHA256 hash of the original text (hex string).
42+
/// The text itself is never stored for security/privacy.
43+
/// </summary>
44+
public required string TextHash { get; init; }
45+
46+
/// <summary>
47+
/// Creates a cache key from the given parameters.
48+
/// The text is hashed using SHA256 and not stored.
49+
/// </summary>
50+
/// <param name="provider">Provider type name.</param>
51+
/// <param name="model">Model name.</param>
52+
/// <param name="vectorDimensions">Vector dimensions.</param>
53+
/// <param name="isNormalized">Whether vectors are normalized.</param>
54+
/// <param name="text">The text to hash.</param>
55+
/// <returns>A new EmbeddingCacheKey instance.</returns>
56+
/// <exception cref="ArgumentNullException">When provider or model is null.</exception>
57+
/// <exception cref="ArgumentOutOfRangeException">When vectorDimensions is less than 1.</exception>
58+
public static EmbeddingCacheKey Create(
59+
string provider,
60+
string model,
61+
int vectorDimensions,
62+
bool isNormalized,
63+
string? text)
64+
{
65+
ArgumentNullException.ThrowIfNull(provider, nameof(provider));
66+
ArgumentNullException.ThrowIfNull(model, nameof(model));
67+
ArgumentOutOfRangeException.ThrowIfLessThan(vectorDimensions, 1, nameof(vectorDimensions));
68+
69+
// Normalize null to empty string for consistent hashing
70+
text ??= string.Empty;
71+
72+
return new EmbeddingCacheKey
73+
{
74+
Provider = provider,
75+
Model = model,
76+
VectorDimensions = vectorDimensions,
77+
IsNormalized = isNormalized,
78+
TextLength = text.Length,
79+
TextHash = ComputeSha256Hash(text)
80+
};
81+
}
82+
83+
/// <summary>
84+
/// Generates a composite key string for use as a database primary key.
85+
/// Format: Provider|Model|Dimensions|IsNormalized|TextLength|TextHash
86+
/// </summary>
87+
/// <returns>A string suitable for use as a cache key.</returns>
88+
public string ToCompositeKey()
89+
{
90+
return $"{this.Provider}|{this.Model}|{this.VectorDimensions}|{this.IsNormalized}|{this.TextLength}|{this.TextHash}";
91+
}
92+
93+
/// <summary>
94+
/// Computes SHA256 hash of the input text and returns as lowercase hex string.
95+
/// </summary>
96+
/// <param name="text">The text to hash.</param>
97+
/// <returns>64-character lowercase hex string.</returns>
98+
private static string ComputeSha256Hash(string text)
99+
{
100+
byte[] bytes = SHA256.HashData(Encoding.UTF8.GetBytes(text));
101+
return Convert.ToHexStringLower(bytes);
102+
}
103+
}
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
using KernelMemory.Core.Config.Enums;
3+
4+
namespace KernelMemory.Core.Embeddings.Cache;
5+
6+
/// <summary>
7+
/// Interface for embedding cache implementations.
8+
/// Supports dependency injection and multiple cache implementations (SQLite, etc.).
9+
/// </summary>
10+
public interface IEmbeddingCache
11+
{
12+
/// <summary>
13+
/// Cache mode (read-write, read-only, write-only).
14+
/// Controls whether read and write operations are allowed.
15+
/// </summary>
16+
CacheModes Mode { get; }
17+
18+
/// <summary>
19+
/// Try to retrieve a cached embedding by key.
20+
/// Returns null if not found or if mode is WriteOnly.
21+
/// </summary>
22+
/// <param name="key">The cache key to look up.</param>
23+
/// <param name="ct">Cancellation token.</param>
24+
/// <returns>The cached embedding if found, null otherwise.</returns>
25+
Task<CachedEmbedding?> TryGetAsync(EmbeddingCacheKey key, CancellationToken ct = default);
26+
27+
/// <summary>
28+
/// Store an embedding in the cache.
29+
/// Does nothing if mode is ReadOnly.
30+
/// </summary>
31+
/// <param name="key">The cache key.</param>
32+
/// <param name="vector">The embedding vector to store.</param>
33+
/// <param name="ct">Cancellation token.</param>
34+
Task StoreAsync(EmbeddingCacheKey key, float[] vector, CancellationToken ct = default);
35+
}

0 commit comments

Comments
 (0)