Skip to content

Commit ed1ac5d

Browse files
alkampfergitalkampfergit
committed
Refactor Anthropic
1 parent a6b26f0 commit ed1ac5d

File tree

6 files changed

+126
-61
lines changed

6 files changed

+126
-61
lines changed

‎src/KernelMemory.Extensions.FunctionalTests/Anthropic/AnthropicTests.cs‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
using KernelMemory.ElasticSearch.Anthropic;
2+
using Microsoft.Extensions.DependencyInjection;
3+
using Microsoft.KernelMemory.AI.Anthropic;
4+
5+
namespace KernelMemory.Extensions.FunctionalTests;
6+
7+
public class AnthropicTests
8+
{
9+
private readonly ServiceProvider _serviceProvider;
10+
11+
public AnthropicTests()
12+
{
13+
var services = new ServiceCollection();
14+
services.AddHttpClient<RawAnthropicHttpClient>()
15+
.AddStandardResilienceHandler(options =>
16+
{
17+
// Configure standard resilience options here
18+
});
19+
20+
var anthropicApiKey = Environment.GetEnvironmentVariable("ANTHROPIC_API_KEY")!;
21+
if (string.IsNullOrEmpty(anthropicApiKey))
22+
{
23+
throw new Exception("ANTHROPIC_API_KEY is not set");
24+
}
25+
26+
var config = new AnthropicTextGenerationConfiguration()
27+
{
28+
ApiKey = anthropicApiKey,
29+
};
30+
services.AddSingleton(config);
31+
32+
_serviceProvider = services.BuildServiceProvider();
33+
}
34+
35+
[Theory]
36+
[InlineData(AnthropicTextGenerationConfiguration.HaikuModelName)]
37+
[InlineData(AnthropicTextGenerationConfiguration.Sonnet35ModelName)]
38+
public async Task Simple_call_some_models(string modelName)
39+
{
40+
var client = new RawAnthropicClient(_serviceProvider);
41+
var parameters = new CallClaudeStreamingParams()
42+
{
43+
ModelName = modelName,
44+
System = "You are an helpful assistant!",
45+
Prompt = "What is the capital of the United States?",
46+
MaxTokens = 5,
47+
};
48+
49+
var result = client.CallClaudeStreaming(parameters);
50+
var response = await result.ToListAsync();
51+
Assert.True(response.Count > 0);
52+
}
53+
}

‎src/KernelMemory.Extensions.FunctionalTests/Cohere/CohereTests.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@ namespace KernelMemory.Extensions.FunctionalTests.Cohere;
88

99
public class CohereTests
1010
{
11-
private ServiceProvider _serviceProvider;
11+
private readonly ServiceProvider _serviceProvider;
1212

13-
private IHttpClientFactory _httpClientFactory;
13+
private readonly IHttpClientFactory _httpClientFactory;
1414

1515
public CohereTests()
1616
{

‎src/KernelMemory.Extensions/Anthropic/AnthropicDependencyInjection.cs‎

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,24 +2,23 @@
22
using Microsoft.KernelMemory;
33
using Microsoft.KernelMemory.AI;
44

5-
namespace KernelMemory.ElasticSearch.Anthropic
5+
namespace KernelMemory.ElasticSearch.Anthropic;
6+
7+
public static class AnthropicDependencyInjection
68
{
7-
public static class AnthropicDependencyInjection
9+
public static IKernelMemoryBuilder WithAnthropicTextGeneration(
10+
this IKernelMemoryBuilder builder,
11+
AnthropicTextGenerationConfiguration config)
812
{
9-
public static IKernelMemoryBuilder WithAnthropicTextGeneration(
10-
this IKernelMemoryBuilder builder,
11-
AnthropicTextGenerationConfiguration config)
12-
{
13-
builder.Services.AddAnthropicTextGeneration(config);
14-
return builder;
15-
}
13+
builder.Services.AddAnthropicTextGeneration(config);
14+
return builder;
15+
}
1616

17-
public static IServiceCollection AddAnthropicTextGeneration(
18-
this IServiceCollection services,
19-
AnthropicTextGenerationConfiguration config)
20-
{
21-
services.AddSingleton(config);
22-
return services.AddSingleton<ITextGenerator, AnthropicTextGeneration>();
23-
}
17+
public static IServiceCollection AddAnthropicTextGeneration(
18+
this IServiceCollection services,
19+
AnthropicTextGenerationConfiguration config)
20+
{
21+
services.AddSingleton(config);
22+
return services.AddSingleton<ITextGenerator, AnthropicTextGeneration>();
2423
}
2524
}

‎src/KernelMemory.Extensions/Anthropic/AnthropicTextGeneration.cs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ namespace KernelMemory.ElasticSearch.Anthropic;
1111

1212
internal class AnthropicTextGeneration : ITextGenerator
1313
{
14+
private readonly RawAnthropicHttpClient _rawAnthropicHttpClient;
1415
private readonly AnthropicTextGenerationConfiguration _config;
15-
private readonly RawAnthropicClient _client;
1616

1717
/// <summary>
1818
/// We do not have cohere tokenizer directly in C# - in this first version we use gpt4 tokenizer
@@ -21,11 +21,11 @@ internal class AnthropicTextGeneration : ITextGenerator
2121
private static readonly Tokenizer _tokenizer = Tokenizer.CreateTiktokenForModel("gpt-4");
2222

2323
public AnthropicTextGeneration(
24-
IHttpClientFactory httpClientFactory,
24+
RawAnthropicHttpClient rawAnthropicHttpClient,
2525
AnthropicTextGenerationConfiguration config)
2626
{
27+
_rawAnthropicHttpClient = rawAnthropicHttpClient;
2728
_config = config;
28-
_client = new RawAnthropicClient(_config.ApiKey, httpClientFactory, _config.HttpClientName);
2929
}
3030

3131
/// <inheritdoc />
@@ -51,7 +51,7 @@ public async IAsyncEnumerable<string> GenerateTextAsync(
5151
Temperature = options.Temperature,
5252
MaxTokens = options.MaxTokens ?? 2048
5353
};
54-
var streamedResponse = _client.CallClaudeStreaming(p);
54+
var streamedResponse = _rawAnthropicHttpClient.CallClaudeStreaming(p);
5555

5656
await foreach (var response in streamedResponse.WithCancellation(cancellationToken))
5757
{

‎src/KernelMemory.Extensions/Anthropic/AnthropicTextGenerationConfiguration.cs‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,16 @@
1-
namespace KernelMemory.ElasticSearch.Anthropic
2-
{
3-
public class AnthropicTextGenerationConfiguration
4-
{
5-
/// <summary>
6-
/// This allows configuring the client that will be used for httpclient factory to create client.
7-
/// </summary>
8-
public string? HttpClientName { get; set; }
1+
namespace KernelMemory.ElasticSearch.Anthropic;
92

10-
public int MaxTokenTotal { get; set; } = 4096;
3+
public class AnthropicTextGenerationConfiguration
4+
{
5+
public int MaxTokenTotal { get; set; } = 4096;
116

12-
public string ApiKey { get; set; }
7+
public string ApiKey { get; set; }
138

14-
public string ModelName { get; set; } = HaikuModelName;
9+
public string ModelName { get; set; } = HaikuModelName;
1510

16-
public const string HaikuModelName = "claude-3-haiku-20240307";
17-
public const string SonnetModelName = "claude-3-sonnet-20240229";
18-
public const string OpusModelName = "claude-3-opus-20240229";
19-
}
11+
public const string HaikuModelName = "claude-3-haiku-20240307";
12+
public const string SonnetModelName = "claude-3-sonnet-20240229";
13+
14+
public const string Sonnet35ModelName = "claude-3-5-sonnet-20240620";
15+
public const string OpusModelName = "claude-3-opus-20240229";
2016
}

‎src/KernelMemory.Extensions/Anthropic/RawAnthropicClient.cs‎

Lines changed: 41 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using KernelMemory.Extensions.Helper;
2+
using Microsoft.Extensions.DependencyInjection;
23
using System;
34
using System.Collections.Generic;
45
using System.IO;
@@ -12,21 +13,46 @@
1213

1314
namespace KernelMemory.ElasticSearch.Anthropic;
1415

15-
public class RawAnthropicClient
16+
public class RawAnthropicClient
1617
{
17-
private readonly string _apiKey;
18-
private readonly IHttpClientFactory _httpClientFactory;
19-
private readonly string? _httpClientName;
20-
private readonly string _baseUrl = "https://api.anthropic.com";
18+
private readonly IServiceProvider _serviceProvider;
2119

2220
public RawAnthropicClient(
23-
string apiKey,
24-
IHttpClientFactory httpClientFactory,
25-
string? httpClientName)
21+
IServiceProvider serviceProvider
22+
)
23+
{
24+
this._serviceProvider = serviceProvider;
25+
}
26+
27+
public async IAsyncEnumerable<StreamingResponseMessage> CallClaudeStreaming(
28+
CallClaudeStreamingParams parameters,
29+
[EnumeratorCancellation] CancellationToken cancellationToken = default)
30+
{
31+
var client = CreateClient();
32+
await foreach (var message in client.CallClaudeStreaming(parameters, cancellationToken))
33+
{
34+
yield return message;
35+
}
36+
}
37+
38+
private RawAnthropicHttpClient CreateClient()
2639
{
27-
_apiKey = apiKey;
28-
_httpClientFactory = httpClientFactory;
29-
_httpClientName = httpClientName;
40+
return _serviceProvider.GetRequiredService<RawAnthropicHttpClient>();
41+
}
42+
}
43+
44+
public class RawAnthropicHttpClient
45+
{
46+
private readonly AnthropicTextGenerationConfiguration _configuration;
47+
private readonly HttpClient _httpClient;
48+
private readonly string _baseUrl = "https://api.anthropic.com";
49+
50+
public RawAnthropicHttpClient(
51+
AnthropicTextGenerationConfiguration configuration,
52+
HttpClient httpClient)
53+
{
54+
this._configuration = configuration;
55+
_httpClient = httpClient;
3056
}
3157

3258
/// <summary>
@@ -60,15 +86,15 @@ public async IAsyncEnumerable<StreamingResponseMessage> CallClaudeStreaming(
6086
string jsonPayload = HttpClientPayloadSerializerHelper.Serialize(requestPayload);
6187
var content = new StringContent(jsonPayload, Encoding.UTF8, "application/json");
6288
content.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue("application/json");
63-
content.Headers.Add("x-api-key", _apiKey);
89+
content.Headers.Add("x-api-key", _configuration.ApiKey);
6490
content.Headers.Add("anthropic-version", "2023-06-01");
6591

6692
var request = new HttpRequestMessage(HttpMethod.Post, $"{_baseUrl}/v1/messages")
6793
{
6894
Content = content,
6995
};
7096

71-
var httpClient = GetHttpClient();
97+
var httpClient = _httpClient;
7298
var response = await httpClient.SendAsync(request, cancellationToken);
7399
if (!response.IsSuccessStatusCode)
74100
{
@@ -118,15 +144,6 @@ public async IAsyncEnumerable<StreamingResponseMessage> CallClaudeStreaming(
118144
}
119145
}
120146

121-
private HttpClient GetHttpClient()
122-
{
123-
if (string.IsNullOrEmpty(_httpClientName))
124-
{
125-
return _httpClientFactory.CreateClient();
126-
}
127-
return _httpClientFactory.CreateClient(_httpClientName);
128-
}
129-
130147
public async Task<MessageResponse> CallClaude(string prompt)
131148
{
132149
var requestPayload = new MessageRequest
@@ -148,15 +165,15 @@ public async Task<MessageResponse> CallClaude(string prompt)
148165
string jsonPayload = HttpClientPayloadSerializerHelper.Serialize(requestPayload);
149166
var content = new StringContent(jsonPayload, Encoding.UTF8, "application/json");
150167
content.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue("application/json");
151-
content.Headers.Add("x-api-key", _apiKey);
168+
content.Headers.Add("x-api-key", _configuration.ApiKey);
152169
content.Headers.Add("anthropic-version", "2023-06-01");
153170

154171
var request = new HttpRequestMessage(HttpMethod.Post, $"{_baseUrl}/v1/messages")
155172
{
156173
Content = content,
157174
};
158175

159-
var httpClient = GetHttpClient();
176+
var httpClient = _httpClient;
160177
var response = await httpClient.SendAsync(request);
161178
if (!response.IsSuccessStatusCode)
162179
{

0 commit comments

Comments
 (0)