Skip to content

Commit df9091c

Browse files
committed
feat: add embedding generators and cache (#7)
Implement embedding generation infrastructure with caching support: Core Components: - IEmbeddingGenerator interface for all providers - EmbeddingConstants with known model dimensions - CachedEmbeddingGenerator decorator for transparent caching Cache System: - IEmbeddingCache interface with SQLite implementation - EmbeddingCacheKey with SHA256 hashing (privacy: text never stored) - CachedEmbedding model with provider/model metadata - CacheModes enum: ReadWrite, ReadOnly, WriteOnly Embedding Providers: - OllamaEmbeddingGenerator for local models - OpenAIEmbeddingGenerator for OpenAI API - AzureOpenAIEmbeddingGenerator for Azure OpenAI - HuggingFaceEmbeddingGenerator for HuggingFace Inference API Configuration: - HuggingFaceEmbeddingsConfig with JSON discriminator - Updated EmbeddingsTypes enum with HuggingFace value Test coverage: 86.33% (133 new tests)
1 parent 0c05ff7 commit df9091c

27 files changed

+4184
-1
lines changed

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: 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.Diagnostics.CodeAnalysis;
3+
4+
namespace KernelMemory.Core.Embeddings.Cache;
5+
6+
/// <summary>
7+
/// Represents a cached embedding vector with metadata.
8+
/// Stores the vector, optional token count, and timestamp of when it was cached.
9+
/// </summary>
10+
public sealed class CachedEmbedding
11+
{
12+
/// <summary>
13+
/// The embedding vector as a float array.
14+
/// Array is intentional for performance - embeddings are read-only after creation.
15+
/// </summary>
16+
[SuppressMessage("Performance", "CA1819:Properties should not return arrays",
17+
Justification = "Embedding vectors are performance-critical and read-only after creation")]
18+
public required float[] Vector { get; init; }
19+
20+
/// <summary>
21+
/// Optional token count from the provider response.
22+
/// Null if the provider did not return token count.
23+
/// </summary>
24+
public int? TokenCount { get; init; }
25+
26+
/// <summary>
27+
/// Timestamp when this embedding was stored in the cache.
28+
/// </summary>
29+
public required DateTimeOffset Timestamp { get; init; }
30+
}
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
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, model, or text 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+
ArgumentNullException.ThrowIfNull(text, nameof(text));
68+
ArgumentOutOfRangeException.ThrowIfLessThan(vectorDimensions, 1, nameof(vectorDimensions));
69+
70+
return new EmbeddingCacheKey
71+
{
72+
Provider = provider,
73+
Model = model,
74+
VectorDimensions = vectorDimensions,
75+
IsNormalized = isNormalized,
76+
TextLength = text.Length,
77+
TextHash = ComputeSha256Hash(text)
78+
};
79+
}
80+
81+
/// <summary>
82+
/// Generates a composite key string for use as a database primary key.
83+
/// Format: Provider|Model|Dimensions|IsNormalized|TextLength|TextHash
84+
/// </summary>
85+
/// <returns>A string suitable for use as a cache key.</returns>
86+
public string ToCompositeKey()
87+
{
88+
return $"{this.Provider}|{this.Model}|{this.VectorDimensions}|{this.IsNormalized}|{this.TextLength}|{this.TextHash}";
89+
}
90+
91+
/// <summary>
92+
/// Computes SHA256 hash of the input text and returns as lowercase hex string.
93+
/// </summary>
94+
/// <param name="text">The text to hash.</param>
95+
/// <returns>64-character lowercase hex string.</returns>
96+
private static string ComputeSha256Hash(string text)
97+
{
98+
byte[] bytes = SHA256.HashData(Encoding.UTF8.GetBytes(text));
99+
return Convert.ToHexStringLower(bytes);
100+
}
101+
}
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
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="tokenCount">Optional token count from the provider.</param>
34+
/// <param name="ct">Cancellation token.</param>
35+
Task StoreAsync(EmbeddingCacheKey key, float[] vector, int? tokenCount, CancellationToken ct = default);
36+
}

0 commit comments

Comments
 (0)