Skip to content

Commit fe5432d

Browse files
alkampfergitalkampfergit
committed
Example of llama cloud integration
1 parent aefc76e commit fe5432d

File tree

4 files changed

+321
-2
lines changed

4 files changed

+321
-2
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ static async Task Main(string[] args)
2222
services.AddSingleton<CustomSearchPipelineBase>();
2323
services.AddSingleton<AnthropicSample>();
2424
services.AddSingleton<ContextualRetrievalSample>();
25+
services.AddSingleton<CustomParsersSample>();
2526
services.AddHttpClient();
2627

2728
var serviceProvider = services.BuildServiceProvider();
@@ -35,6 +36,7 @@ static async Task Main(string[] args)
3536
["Custom Search pipeline (Basic)"] = typeof(CustomSearchPipelineBase),
3637
["Anthropic"] = typeof(AnthropicSample),
3738
["Contextual retrieval"] = typeof(ContextualRetrievalSample),
39+
["Advanced Parsing"] = typeof(CustomParsersSample),
3840
["Exit"] = null
3941
};
4042

@@ -64,6 +66,7 @@ static async Task Main(string[] args)
6466
@"c:\temp\advancedapisecurity.pdf",
6567
@"S:\OneDrive\B19553_11.pdf",
6668
@"c:\temp\blackhatpython.pdf",
69+
@"c:\temp\manualeDreame.pdf",
6770
@"/Users/gianmariaricci/Downloads/llchaindata/blackhatpython.pdf"]));
6871
await sampleInstance1.RunSample(book);
6972
}

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

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
using KernelMemory.Extensions.ConsoleTest.Helper;
2+
using Microsoft.Extensions.DependencyInjection;
3+
using Microsoft.Extensions.Http.Resilience;
4+
using Microsoft.Extensions.Logging;
5+
using Microsoft.KernelMemory;
6+
using Microsoft.KernelMemory.Context;
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+
13+
namespace SemanticMemory.Samples;
14+
15+
internal class CustomParsersSample : ISample
16+
{
17+
public async Task RunSample(string fileToParse)
18+
{
19+
var services = new ServiceCollection();
20+
21+
services.AddLogging(l => l
22+
.SetMinimumLevel(LogLevel.Trace)
23+
.AddConsole()
24+
.AddDebug()
25+
);
26+
27+
var builder = CreateBasicKernelMemoryBuilder(services);
28+
29+
var serviceProvider = services.BuildServiceProvider();
30+
var parserClient = serviceProvider.GetRequiredService<LLamaCloudParserClient>();
31+
32+
//This is not so goot, but it seems that when we build the ServerlessMemory object
33+
//it cannot access the http services registered in the service collection
34+
builder.Services.AddSingleton(parserClient);
35+
36+
var kernelMemory = builder.Build<MemoryServerless>();
37+
38+
var orchestrator = builder.GetOrchestrator();
39+
40+
var decoders = serviceProvider.GetServices<IContentDecoder>();
41+
42+
// Add pipeline handlers
43+
Console.WriteLine("* Defining pipeline handlers...");
44+
45+
TextExtractionHandler textExtraction = new("extract", orchestrator, decoders);
46+
await orchestrator.AddHandlerAsync(textExtraction);
47+
48+
TextPartitioningHandler textPartitioning = new("partition", orchestrator);
49+
await orchestrator.AddHandlerAsync(textPartitioning);
50+
51+
GenerateEmbeddingsHandler textEmbedding = new("gen_embeddings", orchestrator);
52+
await orchestrator.AddHandlerAsync(textEmbedding);
53+
54+
SaveRecordsHandler saveRecords = new("save_records", orchestrator);
55+
await orchestrator.AddHandlerAsync(saveRecords);
56+
57+
var fileName = Path.GetFileName(fileToParse);
58+
59+
var contextProvider = serviceProvider.GetRequiredService<IContextProvider>();
60+
61+
// now we are going to index document, llamacloud can use caching so we can avoid asking for file.
62+
var pipelineBuilder = orchestrator
63+
.PrepareNewDocumentUpload(
64+
index: "llamacloud",
65+
documentId: fileName,
66+
new TagCollection { { "example", "books" } })
67+
.AddUploadFile(fileName, fileName, fileToParse)
68+
.Then("extract")
69+
.Then("partition")
70+
.Then("gen_embeddings")
71+
.Then("save_records");
72+
73+
contextProvider.AddLLamaCloudParserOptions(fileName, "This is a manual for Dreame vacuum cleaner, I need you to extract a series of sections that can be useful for an helpdesk to answer user questions. You will create sections where each sections contains a question and an answer taken from the text.");
74+
75+
var pipeline = pipelineBuilder.Build();
76+
await orchestrator.RunPipelineAsync(pipeline);
77+
78+
// now ask a question to the user continuously until the user ask an empty question
79+
string question;
80+
do
81+
{
82+
Console.WriteLine("Ask a question to the kernel memory:");
83+
question = Console.ReadLine();
84+
if (!string.IsNullOrWhiteSpace(question))
85+
{
86+
var response = await kernelMemory.AskAsync(question);
87+
Console.WriteLine(response.Result);
88+
}
89+
} while (!string.IsNullOrWhiteSpace(question));
90+
}
91+
92+
private static IKernelMemoryBuilder CreateBasicKernelMemoryBuilder(
93+
ServiceCollection services)
94+
{
95+
// we need a series of services to use Kernel Memory, the first one is
96+
// an embedding service that will be used to create dense vector for
97+
// pieces of test. We can use standard ADA embedding service
98+
var embeddingConfig = new AzureOpenAIConfig
99+
{
100+
APIKey = Dotenv.Get("OPENAI_API_KEY"),
101+
Deployment = "text-embedding-ada-002",
102+
Endpoint = Dotenv.Get("AZURE_ENDPOINT"),
103+
APIType = AzureOpenAIConfig.APITypes.EmbeddingGeneration,
104+
Auth = AzureOpenAIConfig.AuthTypes.APIKey
105+
};
106+
107+
// Now kenel memory needs the LLM data to be able to pass question
108+
// and retreived segments to the model. We can Use GPT35
109+
var chatConfig = new AzureOpenAIConfig
110+
{
111+
APIKey = Dotenv.Get("OPENAI_API_KEY"),
112+
Deployment = Dotenv.Get("KERNEL_MEMORY_DEPLOYMENT_NAME"),
113+
Endpoint = Dotenv.Get("AZURE_ENDPOINT"),
114+
APIType = AzureOpenAIConfig.APITypes.ChatCompletion,
115+
Auth = AzureOpenAIConfig.AuthTypes.APIKey,
116+
MaxTokenTotal = 4096
117+
};
118+
119+
var kernelMemoryBuilder = new KernelMemoryBuilder(services)
120+
.WithAzureOpenAITextGeneration(chatConfig)
121+
.WithAzureOpenAITextEmbeddingGeneration(embeddingConfig);
122+
123+
kernelMemoryBuilder
124+
.WithSimpleFileStorage(new SimpleFileStorageConfig()
125+
{
126+
Directory = "c:\\temp\\kmcps\\storage",
127+
StorageType = FileSystemTypes.Disk
128+
})
129+
.WithSimpleVectorDb(new SimpleVectorDbConfig()
130+
{
131+
Directory = "c:\\temp\\kmcps\\vectorstorage",
132+
StorageType = FileSystemTypes.Disk
133+
});
134+
135+
kernelMemoryBuilder.WithContentDecoder<LLamaCloudParserDocumentDecoder>();
136+
137+
var llamaApiKey = Environment.GetEnvironmentVariable("LLAMA_API_KEY");
138+
if (string.IsNullOrEmpty(llamaApiKey))
139+
{
140+
throw new Exception("LLAMA_API_KEY is not set");
141+
}
142+
143+
//Create llamaparser client
144+
services.AddSingleton(new CloudParserConfiguration
145+
{
146+
ApiKey = llamaApiKey,
147+
});
148+
149+
services.AddHttpClient<LLamaCloudParserClient>()
150+
.AddStandardResilienceHandler(options =>
151+
{
152+
// Configure standard resilience options here
153+
options.TotalRequestTimeout = new HttpTimeoutStrategyOptions()
154+
{
155+
Timeout = TimeSpan.FromMinutes(10),
156+
};
157+
});
158+
159+
services.AddSingleton(sp => sp.GetRequiredService<ILoggerFactory>().CreateLogger<LLamaCloudParserClient>());
160+
161+
services.AddSingleton<IKernelMemoryBuilder>(kernelMemoryBuilder);
162+
return kernelMemoryBuilder;
163+
}
164+
}

‎src/KernelMemory.Extensions/llamaindex/LLamaCloudParserClient.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,8 +117,8 @@ public async Task<bool> WaitForJobSuccessAsync(string jobId, TimeSpan timeout)
117117

118118
public class CloudParserConfiguration
119119
{
120-
public string? ApiKey { get; internal set; }
121-
public string? BaseUrl { get; internal set; } = "https://api.cloud.llamaindex.ai";
120+
public string? ApiKey { get; set; }
121+
public string? BaseUrl { get; set; } = "https://api.cloud.llamaindex.ai";
122122
}
123123

124124
/// <summary>

‎src/KernelMemory.Extensions/llamaindex/LLamaCloudParserDocumentDecoder.cs

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
using Microsoft.Extensions.Logging;
2+
using Microsoft.KernelMemory.Context;
3+
using Microsoft.KernelMemory.DataFormats;
4+
using Microsoft.KernelMemory.Diagnostics;
5+
using Microsoft.KernelMemory.Pipeline;
6+
using System;
7+
using System.IO;
8+
using System.Threading;
9+
using System.Threading.Tasks;
10+
11+
public class LLamaCloudParserDocumentDecoder : IContentDecoder
12+
{
13+
private readonly ILogger<LLamaCloudParserDocumentDecoder> _log;
14+
private readonly LLamaCloudParserClient _client;
15+
private readonly IContextProvider _contextProvider;
16+
17+
public LLamaCloudParserDocumentDecoder(
18+
LLamaCloudParserClient client,
19+
IContextProvider contextProvider,
20+
ILoggerFactory? loggerFactory = null)
21+
{
22+
_log = (loggerFactory ?? DefaultLogger.Factory).CreateLogger<LLamaCloudParserDocumentDecoder>();
23+
_client = client;
24+
_contextProvider = contextProvider;
25+
}
26+
27+
/// <inheritdoc />
28+
public bool SupportsMimeType(string mimeType)
29+
{
30+
//Here we can add more mime types
31+
return mimeType != null
32+
&& (
33+
mimeType.StartsWith(MimeTypes.Pdf, StringComparison.OrdinalIgnoreCase)
34+
|| mimeType.StartsWith(MimeTypes.OpenDocumentText, StringComparison.OrdinalIgnoreCase)
35+
);
36+
}
37+
38+
/// <inheritdoc />
39+
public Task<FileContent> DecodeAsync(string filename, CancellationToken cancellationToken = default)
40+
{
41+
using var stream = File.OpenRead(filename);
42+
return this.DecodeAsync(stream, filename, cancellationToken);
43+
}
44+
45+
/// <inheritdoc />
46+
public Task<FileContent> DecodeAsync(BinaryData data, CancellationToken cancellationToken = default)
47+
{
48+
using var stream = data.ToStream();
49+
return DecodeAsync(stream, null, cancellationToken);
50+
}
51+
52+
/// <inheritdoc />
53+
public Task<FileContent> DecodeAsync(Stream data, CancellationToken cancellationToken = default)
54+
{
55+
return DecodeAsync(data, null, cancellationToken);
56+
}
57+
58+
public async Task<FileContent> DecodeAsync(Stream data, string? fileName, CancellationToken cancellationToken = default)
59+
{
60+
_log.LogDebug("Extracting structured text with llamacloud from file");
61+
62+
//retrieve filename and parsing instructions from context
63+
var context = _contextProvider.GetContext();
64+
string parsingInstructions = string.Empty;
65+
if (context.Arguments.TryGetValue(LLamaCloudParserDocumentDecoderExtensions.FileNameKey, out var fileNameContext))
66+
{
67+
fileName = fileNameContext as string ?? string.Empty;
68+
}
69+
if (context.Arguments.TryGetValue(LLamaCloudParserDocumentDecoderExtensions.ParsingInstructionsKey, out var parsingInstructionsContext))
70+
{
71+
parsingInstructions = parsingInstructionsContext as string ?? string.Empty;
72+
}
73+
74+
// ok we need a way to find the correct instruction for the file, so we can use a different configuration
75+
// for each file that we are going to parse.
76+
var parameters = new UploadParameters();
77+
78+
//file name must not be null
79+
if (string.IsNullOrEmpty(fileName))
80+
{
81+
throw new Exception("LLAMA Cloud error: file name is missing");
82+
}
83+
84+
//ok we need a temporary file name that we need to use to upload the file, we need a seekable stream
85+
var tempFileName = Path.Combine(Path.GetTempPath(), Guid.NewGuid().ToString() + "_" + fileName);
86+
using (var writeFileStream = File.Create(tempFileName))
87+
{
88+
await data.CopyToAsync(writeFileStream, cancellationToken);
89+
}
90+
91+
parameters.WithParsingInstructions(parsingInstructions);
92+
93+
var response = await PerformCall(fileName, tempFileName, parameters);
94+
95+
if (response != null && response.ErrorCode != null)
96+
{
97+
throw new Exception($"LLAMA Cloud error: {response.ErrorCode} - {response.ErrorMessage}");
98+
}
99+
100+
if (response == null)
101+
{
102+
throw new Exception("LLAMA Cloud error: no response");
103+
}
104+
105+
var jobId = response.Id.ToString();
106+
107+
//now wait for the job to be completed
108+
var jobResponse = await _client.WaitForJobSuccessAsync(jobId, TimeSpan.FromMinutes(5));
109+
110+
if (!jobResponse)
111+
{
112+
throw new Exception("LLAMA Cloud error: job not completed");
113+
}
114+
115+
// ok now the job is completed, we can get the markdown
116+
var markdown = await _client.GetJobRawMarkdownAsync(jobId);
117+
118+
var result = new FileContent(MimeTypes.MarkDown);
119+
result.Sections.Add(new FileSection(1, markdown, false));
120+
121+
return result;
122+
}
123+
124+
private async Task<UploadResponse> PerformCall(
125+
string fileName,
126+
string physicalTempFileName,
127+
UploadParameters parameters)
128+
{
129+
try
130+
{
131+
await using var tempFileStream = File.OpenRead(physicalTempFileName);
132+
var uploadResponse = await _client.UploadAsync(tempFileStream, fileName, parameters);
133+
return uploadResponse;
134+
}
135+
finally
136+
{
137+
File.Delete(physicalTempFileName);
138+
}
139+
}
140+
}
141+
142+
public static class LLamaCloudParserDocumentDecoderExtensions
143+
{
144+
internal const string FileNameKey = "llamacloud.filename";
145+
internal const string ParsingInstructionsKey = "llamacloud.parsing_instructions";
146+
147+
public static void AddLLamaCloudParserOptions(this IContextProvider contextProvider, string filename, string parseInstructions)
148+
{
149+
contextProvider.SetContextArg(FileNameKey, filename);
150+
contextProvider.SetContextArg(ParsingInstructionsKey, parseInstructions);
151+
}
152+
}

0 commit comments

Comments
 (0)