Skip to content

Commit 797dd62

Browse files
alkampfergitalkampfergit
committed
Update claude cache sample
1 parent 477435a commit 797dd62

File tree

4 files changed

+287
-2
lines changed

4 files changed

+287
-2
lines changed

‎src/KernelMemory.Extensions.ConsoleTest/Program.cs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ static async Task Main(string[] args)
2121
services.AddSingleton<TextCleanerHandler>();
2222
services.AddSingleton<CustomSearchPipelineBase>();
2323
services.AddSingleton<AnthropicSample>();
24+
services.AddSingleton<ContextualRetrievalSample>();
2425
services.AddHttpClient();
2526

2627
var serviceProvider = services.BuildServiceProvider();
@@ -33,6 +34,7 @@ static async Task Main(string[] args)
3334
["SBert in action"] = typeof(SBertSample),
3435
["Custom Search pipeline (Basic)"] = typeof(CustomSearchPipelineBase),
3536
["Anthropic"] = typeof(AnthropicSample),
37+
["Contextual retrieval"] = typeof(ContextualRetrievalSample),
3638
["Exit"] = null
3739
};
3840

@@ -58,7 +60,10 @@ static async Task Main(string[] args)
5860
{
5961
var book = AnsiConsole.Prompt(new SelectionPrompt<string>()
6062
.Title("Select the [green]book[/] to index")
61-
.AddChoices([@"c:\temp\advancedapisecurity.pdf", @"S:\OneDrive\B19553_11.pdf"]));
63+
.AddChoices([
64+
@"c:\temp\advancedapisecurity.pdf",
65+
@"S:\OneDrive\B19553_11.pdf",
66+
@"/Users/gianmariaricci/Downloads/llchaindata/blackhatpython.pdf"]));
6267
await sampleInstance1.RunSample(book);
6368
}
6469
}

‎src/KernelMemory.Extensions.ConsoleTest/Samples/ContextualRetrievalSample.cs‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
using KernelMemory.ElasticSearch.Anthropic;
2+
using KernelMemory.Extensions.Cohere;
3+
using KernelMemory.Extensions.ConsoleTest.Helper;
4+
using Microsoft.Extensions.DependencyInjection;
5+
using Microsoft.Extensions.Logging;
6+
using Microsoft.KernelMemory;
7+
using Microsoft.KernelMemory.DataFormats;
8+
using Microsoft.KernelMemory.DocumentStorage.DevTools;
9+
using Microsoft.KernelMemory.FileSystem.DevTools;
10+
using Microsoft.KernelMemory.Handlers;
11+
using Microsoft.KernelMemory.MemoryStorage.DevTools;
12+
using Spectre.Console;
13+
14+
namespace SemanticMemory.Samples;
15+
16+
public class ContextualRetrievalSample : ISample
17+
{
18+
public async Task RunSample(string bookPdf)
19+
{
20+
var services = new ServiceCollection();
21+
services.AddLogging(l => l
22+
.SetMinimumLevel(LogLevel.Trace)
23+
.AddConsole()
24+
.AddDebug()
25+
);
26+
//do not forget to add decoders
27+
services.AddDefaultContentDecoders();
28+
services.AddHttpClient<RawAnthropicHttpClient>()
29+
.AddStandardResilienceHandler(options =>
30+
{
31+
// Configure standard resilience options here
32+
});
33+
services.AddTransient<RawAnthropicClient>();
34+
35+
var anthropicApiKey = Environment.GetEnvironmentVariable("ANTHROPIC_API_KEY")!;
36+
if (string.IsNullOrEmpty(anthropicApiKey))
37+
{
38+
throw new Exception("ANTHROPIC_API_KEY is not set");
39+
}
40+
41+
var config = new AnthropicTextGenerationConfiguration()
42+
{
43+
ApiKey = anthropicApiKey,
44+
};
45+
services.AddSingleton(config);
46+
47+
48+
var builder = CreateBasicKernelMemoryBuilder(services);
49+
var kernelMemory = builder.Build<MemoryServerless>();
50+
51+
var orchestrator = builder.GetOrchestrator();
52+
53+
var serviceProvider = services.BuildServiceProvider();
54+
var decoders = serviceProvider.GetServices<IContentDecoder>();
55+
56+
var anthropicClient = serviceProvider.GetRequiredService<RawAnthropicClient>();
57+
// Add pipeline handlers
58+
Console.WriteLine("* Defining pipeline handlers...");
59+
60+
TextExtractionHandler textExtraction = new("extract", orchestrator, decoders);
61+
await orchestrator.AddHandlerAsync(textExtraction);
62+
63+
TextCleanerHandler textCleanerHandler = new("clean", orchestrator);
64+
await orchestrator.AddHandlerAsync(textCleanerHandler);
65+
66+
TextPartitioningHandler textPartitioning = new("partition", orchestrator);
67+
await orchestrator.AddHandlerAsync(textPartitioning);
68+
69+
ClaudeContextualRetrievalHandler contextual = new("contextual", orchestrator, anthropicClient);
70+
await orchestrator.AddHandlerAsync(contextual);
71+
72+
GenerateEmbeddingsHandler textEmbedding = new("gen_embeddings", orchestrator);
73+
await orchestrator.AddHandlerAsync(textEmbedding);
74+
75+
SaveRecordsHandler saveRecords = new("save_records", orchestrator);
76+
await orchestrator.AddHandlerAsync(saveRecords);
77+
78+
// orchestrator.AddHandlerAsync(...);
79+
// orchestrator.AddHandlerAsync(...);
80+
81+
// Create sample pipeline with 4 files
82+
var index = AnsiConsole.Confirm("Do you want to index document?");
83+
if (index)
84+
{
85+
Console.WriteLine("Single pipeline for single document");
86+
var pipeline = orchestrator
87+
.PrepareNewDocumentUpload(
88+
index: "book",
89+
documentId: "grayhatpython",
90+
new TagCollection { { "security", "books" } })
91+
.AddUploadFile("file1", Path.GetFileName(bookPdf), bookPdf)
92+
.Then("extract")
93+
.Then("clean")
94+
.Then("partition")
95+
.Then("contextual")
96+
.Then("gen_embeddings")
97+
.Then("save_records")
98+
.Build();
99+
100+
// Execute pipeline
101+
Console.WriteLine("* Executing pipeline...");
102+
await orchestrator.RunPipelineAsync(pipeline);
103+
}
104+
string question;
105+
do
106+
{
107+
Console.WriteLine("Ask a question to the kernel memory:");
108+
question = Console.ReadLine();
109+
if (!string.IsNullOrWhiteSpace(question))
110+
{
111+
var response = await kernelMemory.AskAsync(question, index: "book");
112+
Console.WriteLine(response.Result);
113+
}
114+
} while (!string.IsNullOrWhiteSpace(question));
115+
}
116+
117+
private static IKernelMemoryBuilder CreateBasicKernelMemoryBuilder(
118+
ServiceCollection services)
119+
{
120+
// we need a series of services to use Kernel Memory, the first one is
121+
// an embedding service that will be used to create dense vector for
122+
// pieces of test. We can use standard ADA embedding service
123+
var embeddingConfig = new AzureOpenAIConfig
124+
{
125+
APIKey = Dotenv.Get("OPENAI_API_KEY"),
126+
Deployment = "text-embedding-ada-002",
127+
Endpoint = Dotenv.Get("AZURE_ENDPOINT"),
128+
APIType = AzureOpenAIConfig.APITypes.EmbeddingGeneration,
129+
Auth = AzureOpenAIConfig.AuthTypes.APIKey
130+
};
131+
132+
// Now kenel memory needs the LLM data to be able to pass question
133+
// and retreived segments to the model. We can Use GPT35
134+
var chatConfig = new AzureOpenAIConfig
135+
{
136+
APIKey = Dotenv.Get("OPENAI_API_KEY"),
137+
Deployment = Dotenv.Get("KERNEL_MEMORY_DEPLOYMENT_NAME"),
138+
Endpoint = Dotenv.Get("AZURE_ENDPOINT"),
139+
APIType = AzureOpenAIConfig.APITypes.ChatCompletion,
140+
Auth = AzureOpenAIConfig.AuthTypes.APIKey,
141+
MaxTokenTotal = 4096
142+
};
143+
144+
var kernelMemoryBuilder = new KernelMemoryBuilder(services)
145+
.WithAzureOpenAITextGeneration(chatConfig)
146+
.WithAzureOpenAITextEmbeddingGeneration(embeddingConfig);
147+
148+
kernelMemoryBuilder
149+
.WithSimpleFileStorage(new SimpleFileStorageConfig()
150+
{
151+
Directory = "/tmp/km/storage",
152+
//Directory = "c:\\temp\\km2\\storage",
153+
StorageType = FileSystemTypes.Disk
154+
})
155+
.WithSimpleVectorDb(new SimpleVectorDbConfig()
156+
{
157+
Directory = "/tmp/km/vectorstorage",
158+
//Directory = "c:\\temp\\km2\\vectorstorage",
159+
StorageType = FileSystemTypes.Disk
160+
});
161+
162+
services.AddSingleton<IKernelMemoryBuilder>(kernelMemoryBuilder);
163+
return kernelMemoryBuilder;
164+
}
165+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ public class AnthropicTextGenerationConfiguration
88

99
public string ModelName { get; set; } = HaikuModelName;
1010

11-
public const string HaikuModelName = "claude-3-haiku-20240307";
11+
public const string HaikuModelName = "claude-3-5-haiku-latest";
1212
public const string SonnetModelName = "claude-3-sonnet-20240229";
1313

1414
public const string Sonnet35ModelName = "claude-3-5-sonnet-20240620";

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

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
using KernelMemory.ElasticSearch.Anthropic;
2+
using Microsoft.Extensions.Logging;
3+
using Microsoft.KernelMemory.AI.Anthropic;
4+
using Microsoft.KernelMemory.Diagnostics;
5+
using Microsoft.KernelMemory.Pipeline;
6+
using System.Collections.Generic;
7+
using System.Linq;
8+
using System.Text;
9+
using System.Threading;
10+
using System.Threading.Tasks;
11+
12+
namespace KernelMemory.Extensions.ConsoleTest.Helper;
13+
14+
public class ClaudeContextualRetrievalHandler : IPipelineStepHandler
15+
{
16+
private string _name;
17+
private IPipelineOrchestrator _orchestrator;
18+
private readonly RawAnthropicClient _rawAnthropicHttpClient;
19+
private readonly ILogger<ClaudeContextualRetrievalHandler> _log;
20+
21+
public ClaudeContextualRetrievalHandler(
22+
string name,
23+
IPipelineOrchestrator orchestrator,
24+
RawAnthropicClient rawAnthropicHttpClient,
25+
ILogger<ClaudeContextualRetrievalHandler>? log = null)
26+
{
27+
_name = name;
28+
_orchestrator = orchestrator;
29+
_rawAnthropicHttpClient = rawAnthropicHttpClient;
30+
_log = log ?? DefaultLogger<ClaudeContextualRetrievalHandler>.Instance; ;
31+
}
32+
33+
public string StepName => _name;
34+
35+
public async Task<(bool success, DataPipeline updatedPipeline)> InvokeAsync(DataPipeline pipeline, CancellationToken cancellationToken = default)
36+
{
37+
_log.LogDebug("Partitioning text, pipeline '{0}/{1}'", pipeline.Index, pipeline.DocumentId);
38+
39+
if (pipeline.Files.Count == 0)
40+
{
41+
_log.LogWarning("Pipeline '{0}/{1}': there are no files to process, moving to next pipeline step.", pipeline.Index, pipeline.DocumentId);
42+
return (true, pipeline);
43+
}
44+
45+
foreach (DataPipeline.FileDetails uploadedFile in pipeline.Files)
46+
{
47+
// Track new files being generated (cannot edit originalFile.GeneratedFiles while looping it)
48+
Dictionary<string, DataPipeline.GeneratedFileDetails> newFiles = new();
49+
// we need to find the original extracted text
50+
var textFile = uploadedFile.GeneratedFiles.FirstOrDefault(f => f.Value.ArtifactType == DataPipeline.ArtifactTypes.ExtractedText);
51+
52+
if (textFile.Value != null)
53+
{
54+
var contentFile = await _orchestrator.ReadFileAsync(pipeline, textFile.Value.Name, cancellationToken).ConfigureAwait(false);
55+
var text = contentFile.ToString();
56+
foreach (KeyValuePair<string, DataPipeline.GeneratedFileDetails> generatedFile in uploadedFile.GeneratedFiles)
57+
{
58+
var file = generatedFile.Value;
59+
if (file.AlreadyProcessedBy(this))
60+
{
61+
_log.LogTrace("File {0} already processed by this handler", file.Name);
62+
continue;
63+
}
64+
65+
if (file.ArtifactType != DataPipeline.ArtifactTypes.TextPartition)
66+
{
67+
_log.LogTrace("Skipping file {0} (not text partition)", file.Name);
68+
continue;
69+
}
70+
71+
var content = await _orchestrator.ReadFileAsync(pipeline, file.Name, cancellationToken).ConfigureAwait(false);
72+
if (content.ToArray().Length == 0) { continue; }
73+
74+
var segmentContent = content.ToString();
75+
var contextGenerated = await EnrichTextAsync(text, segmentContent, cancellationToken);
76+
var enrichedText = contextGenerated + "\n" + segmentContent;
77+
await _orchestrator.WriteTextFileAsync(pipeline, file.Name, enrichedText, cancellationToken).ConfigureAwait(false);
78+
79+
file.MarkProcessedBy(this);
80+
}
81+
82+
// Add new files to pipeline status
83+
foreach (var file in newFiles)
84+
{
85+
uploadedFile.GeneratedFiles.Add(file.Key, file.Value);
86+
}
87+
}
88+
}
89+
return (true, pipeline);
90+
}
91+
92+
private async Task<string> EnrichTextAsync(
93+
string text,
94+
string segment,
95+
CancellationToken cancellationToken)
96+
{
97+
var parameters = new CallClaudeStreamingParams()
98+
{
99+
ModelName = AnthropicTextGenerationConfiguration.HaikuModelName,
100+
System = [
101+
SystemMessage.Create(@$"<document>
102+
{text}
103+
</document>", CacheControl.Ephemeral)],
104+
Messages = [Message.Create("user", @$"Here is the chunk we want to situate within the whole document
105+
<chunk>
106+
{segment}
107+
</chunk>
108+
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else. ")],
109+
MaxTokens = 5,
110+
};
111+
112+
var result = await _rawAnthropicHttpClient.CallClaudeAsync(parameters, cancellationToken);
113+
return result.Content[0].Text;
114+
}
115+
}

0 commit comments

Comments
 (0)