Skip to content

Commit 35e3911

Browse files
eerhardteerhardt
authored
Fix Ollama OpenTelemetry pipeline (#509)
* Fix Ollama OpenTelemetry pipeline 1. Enable logging in the OpenTelemetry clients 2. Ensure the OpenTelemetry clients are immediately wrapped around the raw api client. This ensures all calls to the Ollama api are logged correctly. * Fix comment
1 parent 0076c96 commit 35e3911

File tree

4 files changed

+94
-19
lines changed

4 files changed

+94
-19
lines changed

‎src/CommunityToolkit.Aspire.OllamaSharp/AspireOllamaChatClientExtensions.cs‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using Microsoft.Extensions.AI;
22
using Microsoft.Extensions.DependencyInjection;
3+
using Microsoft.Extensions.Logging;
34
using OllamaSharp;
45

56
namespace Microsoft.Extensions.Hosting;
@@ -19,8 +20,7 @@ public static ChatClientBuilder AddChatClient(this AspireOllamaApiClientBuilder
1920
ArgumentNullException.ThrowIfNull(builder, nameof(builder));
2021

2122
return builder.HostBuilder.Services.AddChatClient(
22-
services => CreateInnerChatClient(services, builder))
23-
.UseAspireDefaults();
23+
services => CreateInnerChatClient(services, builder));
2424
}
2525

2626
/// <summary>
@@ -36,10 +36,14 @@ public static ChatClientBuilder AddKeyedChatClient(
3636

3737
return builder.HostBuilder.Services.AddKeyedChatClient(
3838
builder.ServiceKey,
39-
services => CreateInnerChatClient(services, builder))
40-
.UseAspireDefaults();
39+
services => CreateInnerChatClient(services, builder));
4140
}
4241

42+
/// <summary>
43+
/// Wrap the <see cref="IOllamaApiClient"/> in a telemetry client if tracing is enabled.
44+
/// Note that this doesn't use ".UseOpenTelemetry()" because the order of the clients would be incorrect.
45+
/// We want the telemetry client to be the innermost client, right next to the inner <see cref="IOllamaApiClient"/>.
46+
/// </summary>
4347
private static IChatClient CreateInnerChatClient(
4448
IServiceProvider services,
4549
AspireOllamaApiClientBuilder builder)
@@ -48,10 +52,12 @@ private static IChatClient CreateInnerChatClient(
4852

4953
var result = (IChatClient)ollamaApiClient;
5054

51-
return builder.DisableTracing ? result : new OpenTelemetryChatClient(result);
52-
}
55+
if (builder.DisableTracing)
56+
{
57+
return result;
58+
}
5359

54-
private static ChatClientBuilder UseAspireDefaults(this ChatClientBuilder builder) =>
55-
builder.UseLogging()
56-
.UseOpenTelemetry();
60+
var loggerFactory = services.GetService<ILoggerFactory>();
61+
return new OpenTelemetryChatClient(result, loggerFactory?.CreateLogger(typeof(OpenTelemetryChatClient)));
62+
}
5763
}

‎src/CommunityToolkit.Aspire.OllamaSharp/AspireOllamaEmbeddingGeneratorExtensions.cs‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using Microsoft.Extensions.AI;
22
using Microsoft.Extensions.DependencyInjection;
3+
using Microsoft.Extensions.Logging;
34
using OllamaSharp;
45

56
namespace Microsoft.Extensions.Hosting;
@@ -18,7 +19,7 @@ public static EmbeddingGeneratorBuilder<string, Embedding<float>> AddEmbeddingGe
1819
this AspireOllamaApiClientBuilder builder)
1920
{
2021
return builder.HostBuilder.Services.AddEmbeddingGenerator(
21-
services => CreateInnerEmbeddingGenerator(services, builder)).UseAspireDefaults();
22+
services => CreateInnerEmbeddingGenerator(services, builder));
2223
}
2324

2425
/// <summary>
@@ -34,9 +35,14 @@ public static EmbeddingGeneratorBuilder<string, Embedding<float>> AddKeyedEmbedd
3435

3536
return builder.HostBuilder.Services.AddKeyedEmbeddingGenerator(
3637
builder.ServiceKey,
37-
services => CreateInnerEmbeddingGenerator(services, builder)).UseAspireDefaults();
38+
services => CreateInnerEmbeddingGenerator(services, builder));
3839
}
3940

41+
/// <summary>
42+
/// Wrap the <see cref="IOllamaApiClient"/> in a telemetry client if tracing is enabled.
43+
/// Note that this doesn't use ".UseOpenTelemetry()" because the order of the clients would be incorrect.
44+
/// We want the telemetry client to be the innermost client, right next to the inner <see cref="IOllamaApiClient"/>.
45+
/// </summary>
4046
private static IEmbeddingGenerator<string, Embedding<float>> CreateInnerEmbeddingGenerator(
4147
IServiceProvider services,
4248
AspireOllamaApiClientBuilder builder)
@@ -45,13 +51,14 @@ private static IEmbeddingGenerator<string, Embedding<float>> CreateInnerEmbeddin
4551

4652
var result = (IEmbeddingGenerator<string, Embedding<float>>)ollamaApiClient;
4753

48-
return builder.DisableTracing
49-
? result
50-
: new OpenTelemetryEmbeddingGenerator<string, Embedding<float>>(result);
51-
}
54+
if (builder.DisableTracing)
55+
{
56+
return result;
57+
}
5258

53-
private static EmbeddingGeneratorBuilder<TKey, TEmbedding> UseAspireDefaults<TKey, TEmbedding>(this EmbeddingGeneratorBuilder<TKey, TEmbedding> builder)
54-
where TEmbedding : Embedding =>
55-
builder.UseLogging()
56-
.UseOpenTelemetry();
59+
var loggerFactory = services.GetService<ILoggerFactory>();
60+
return new OpenTelemetryEmbeddingGenerator<string, Embedding<float>>(
61+
result,
62+
loggerFactory?.CreateLogger(typeof(OpenTelemetryEmbeddingGenerator<string, Embedding<float>>)));
63+
}
5764
}

‎tests/CommunityToolkit.Aspire.OllamaSharp.Tests/OllamaSharpIChatClientTests.cs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
using Microsoft.Extensions.AI;
22
using Microsoft.Extensions.Configuration;
33
using Microsoft.Extensions.Hosting;
4+
using OllamaSharp;
5+
using System.Runtime.CompilerServices;
46

57
namespace CommunityToolkit.Aspire.OllamaSharp.Tests;
68

@@ -122,4 +124,32 @@ public void CanSetMultipleKeyedClients()
122124
Assert.NotEqual(client, client2);
123125
Assert.NotEqual(client, client3);
124126
}
127+
128+
[Fact]
129+
public void CanChainUseMethodsCorrectly()
130+
{
131+
var builder = Host.CreateEmptyApplicationBuilder(null);
132+
builder.Configuration.AddInMemoryCollection([
133+
new KeyValuePair<string, string?>("ConnectionStrings:Ollama",$"Endpoint={Endpoint}")
134+
]);
135+
136+
builder.Services.AddDistributedMemoryCache();
137+
138+
builder.AddOllamaApiClient("Ollama")
139+
.AddChatClient()
140+
.UseDistributedCache()
141+
.UseFunctionInvocation();
142+
143+
using var host = builder.Build();
144+
var client = host.Services.GetRequiredService<IChatClient>();
145+
146+
var distributedCacheClient = Assert.IsType<DistributedCachingChatClient>(client);
147+
var functionInvocationClient = Assert.IsType<FunctionInvokingChatClient>(GetInnerClient(distributedCacheClient));
148+
var otelClient = Assert.IsType<OpenTelemetryChatClient>(GetInnerClient(functionInvocationClient));
149+
150+
Assert.IsType<IOllamaApiClient>(GetInnerClient(otelClient), exactMatch: false);
151+
}
152+
153+
[UnsafeAccessor(UnsafeAccessorKind.Method, Name = "get_InnerClient")]
154+
private static extern IChatClient GetInnerClient(DelegatingChatClient client);
125155
}

‎tests/CommunityToolkit.Aspire.OllamaSharp.Tests/OllamaSharpIEmbeddingGeneratorTests.cs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
using Microsoft.Extensions.AI;
22
using Microsoft.Extensions.Configuration;
33
using Microsoft.Extensions.Hosting;
4+
using OllamaSharp;
5+
using System.Reflection;
6+
using System.Runtime.CompilerServices;
47

58
namespace CommunityToolkit.Aspire.OllamaSharp.Tests;
69

@@ -122,4 +125,33 @@ public void CanSetMultipleKeyedClients()
122125
Assert.NotEqual(client, client2);
123126
Assert.NotEqual(client, client3);
124127
}
128+
129+
[Fact]
130+
public void CanChainUseMethodsCorrectly()
131+
{
132+
var builder = Host.CreateEmptyApplicationBuilder(null);
133+
builder.Configuration.AddInMemoryCollection([
134+
new KeyValuePair<string, string?>("ConnectionStrings:Ollama",$"Endpoint={Endpoint}")
135+
]);
136+
137+
builder.Services.AddDistributedMemoryCache();
138+
139+
builder.AddOllamaApiClient("Ollama")
140+
.AddEmbeddingGenerator()
141+
.UseDistributedCache();
142+
143+
using var host = builder.Build();
144+
var client = host.Services.GetRequiredService<IEmbeddingGenerator<string, Embedding<float>>>();
145+
146+
var distributedCacheClient = Assert.IsType<DistributedCachingEmbeddingGenerator<string, Embedding<float>>>(client);
147+
var otelClient = Assert.IsType<OpenTelemetryEmbeddingGenerator<string, Embedding<float>>>(GetInnerGenerator(distributedCacheClient));
148+
149+
Assert.IsType<IOllamaApiClient>(GetInnerGenerator(otelClient), exactMatch: false);
150+
}
151+
152+
private static IEmbeddingGenerator<TInput, TEmbedding> GetInnerGenerator<TInput, TEmbedding>(DelegatingEmbeddingGenerator<TInput, TEmbedding> generator)
153+
where TEmbedding : Embedding =>
154+
(IEmbeddingGenerator<TInput,TEmbedding>)(generator.GetType()
155+
.GetProperty("InnerGenerator", BindingFlags.Instance | BindingFlags.NonPublic)?
156+
.GetValue(generator, null) ?? throw new InvalidOperationException());
125157
}

0 commit comments

Comments
 (0)