-
Notifications
You must be signed in to change notification settings - Fork 57
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #111 from hsavran/main
PostgreSQL and Generate vector for Mongo DB
- Loading branch information
Showing
23 changed files
with
774 additions
and
12 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
26 changes: 26 additions & 0 deletions
26
.../Cosmos.DataTransfer.MongoVectorExtension/Cosmos.DataTransfer.MongoVectorExtension.csproj
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,26 @@ | ||
| <Project Sdk="Microsoft.NET.Sdk"> | ||
|
|
||
| <PropertyGroup> | ||
| <TargetFramework>net6.0</TargetFramework> | ||
| <ImplicitUsings>enable</ImplicitUsings> | ||
| <Nullable>enable</Nullable> | ||
| <OutputType>Exe</OutputType> | ||
| </PropertyGroup> | ||
|
|
||
| <ItemGroup> | ||
| <PackageReference Include="Azure.AI.OpenAI" Version="1.0.0-beta.12" /> | ||
| <PackageReference Include="Microsoft.Extensions.Configuration.Binder" Version="6.0.0" /> | ||
| <PackageReference Include="MongoDB.Driver" Version="2.19.1" /> | ||
| <PackageReference Include="System.ComponentModel.Composition" Version="6.0.0" /> | ||
| </ItemGroup> | ||
|
|
||
| <ItemGroup> | ||
| <ProjectReference Include="..\..\..\Interfaces\Cosmos.DataTransfer.Interfaces\Cosmos.DataTransfer.Interfaces.csproj" /> | ||
| <ProjectReference Include="..\Cosmos.DataTransfer.MongoExtension\Cosmos.DataTransfer.MongoExtension.csproj" /> | ||
| </ItemGroup> | ||
|
|
||
| <Target Name="PublishToExtensionsFolder" AfterTargets="Build" Condition=" '$(Configuration)' == 'Debug' "> | ||
| <Exec Command="dotnet publish --configuration $(Configuration) --no-build -p:PublishProfile=PublishToExtensionsFolder" /> | ||
| </Target> | ||
|
|
||
| </Project> |
89 changes: 89 additions & 0 deletions
89
Extensions/Mongo/Cosmos.DataTransfer.MongoVectorExtension/MongoVectorDataSinkExtension.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,89 @@ | ||
| using System.ComponentModel.Composition; | ||
| using Azure; | ||
| using Azure.AI.OpenAI; | ||
| using Cosmos.DataTransfer.Interfaces; | ||
| using Cosmos.DataTransfer.MongoExtension; | ||
| using Cosmos.DataTransfer.MongoVectorExtension.Settings; | ||
| using Microsoft.Extensions.Configuration; | ||
| using Microsoft.Extensions.Logging; | ||
| using MongoDB.Bson; | ||
|
|
||
| namespace Cosmos.DataTransfer.MongoVectorExtension; | ||
| [Export(typeof(IDataSinkExtension))] | ||
| public class MongoVectorDataSinkExtension : IDataSinkExtensionWithSettings | ||
| { | ||
| public string DisplayName => $"MongoDB-Vector{ExtensionExtensions.BetaExtensionTag}"; | ||
|
|
||
| public async Task WriteAsync(IAsyncEnumerable<IDataItem> dataItems, IConfiguration config, IDataSourceExtension dataSource, ILogger logger, CancellationToken cancellationToken = default) | ||
| { | ||
| var settings = config.Get<MongoVectorSinkSettings>(); | ||
| settings.Validate(); | ||
|
|
||
| if (!string.IsNullOrEmpty(settings.ConnectionString) && !string.IsNullOrEmpty(settings.DatabaseName) && !string.IsNullOrEmpty(settings.Collection)) | ||
| { | ||
| var Isembeddingsetsvalid = false; | ||
| var client = new OpenAIClient(""); | ||
| if (settings.GenerateEmbedding.HasValue && settings.GenerateEmbedding.Value && settings.SourcePropEmbedding != null && settings.DestPropEmbedding != null) | ||
| { | ||
| if (!string.IsNullOrEmpty(settings.OpenAIUrl) && !string.IsNullOrEmpty(settings.OpenAIKey) && !string.IsNullOrEmpty(settings.OpenAIDeploymentName)) | ||
| { | ||
| client = new OpenAIClient(new Uri(settings.OpenAIUrl), new AzureKeyCredential(settings.OpenAIKey)); | ||
| Isembeddingsetsvalid = true; | ||
| logger.LogInformation("OpenAI Embedding settings are valid."); | ||
| } | ||
| } | ||
|
|
||
| var context = new Context(settings.ConnectionString, settings.DatabaseName); | ||
| var repo = context.GetRepository<BsonDocument>(settings.Collection); | ||
| var batchSize = settings.BatchSize ?? 1000; | ||
| var objects = new List<BsonDocument>(); | ||
| int itemCount = 0; | ||
| await foreach (var item in dataItems.WithCancellation(cancellationToken)) | ||
| { | ||
| var dict = item.BuildDynamicObjectTree(); | ||
|
|
||
| if (Isembeddingsetsvalid) | ||
| { | ||
| var valtoemb = item.GetValue(settings.SourcePropEmbedding)?.ToString(); | ||
| if (!string.IsNullOrEmpty(valtoemb) && valtoemb?.Length < 8192) | ||
| { | ||
| var options = new EmbeddingsOptions() | ||
| { | ||
| DeploymentName = settings.OpenAIDeploymentName, | ||
| Input = { valtoemb } | ||
| }; | ||
| var vector = await client.GetEmbeddingsAsync(options,cancellationToken); | ||
| if (vector != null) | ||
| { | ||
| dict?.TryAdd(settings.DestPropEmbedding, vector.Value.Data[0].Embedding.ToArray()); | ||
| } | ||
| } | ||
| } | ||
| objects.Add(new BsonDocument(dict)); | ||
| itemCount++; | ||
|
|
||
| if (objects.Count == batchSize) | ||
| { | ||
| await repo.AddRange(objects); | ||
| logger.LogInformation("Added {ItemCount} items to collection '{Collection}'", itemCount, settings.Collection); | ||
| objects.Clear(); | ||
| } | ||
| } | ||
|
|
||
| if (objects.Any()) | ||
| { | ||
| await repo.AddRange(objects); | ||
| } | ||
|
|
||
| if (itemCount > 0) | ||
| logger.LogInformation("Added {ItemCount} total items to collection '{Collection}'", itemCount, settings.Collection); | ||
| else | ||
| logger.LogWarning("No items added to collection '{Collection}'", settings.Collection); | ||
| } | ||
| } | ||
|
|
||
| public IEnumerable<IDataExtensionSettings> GetSettings() | ||
| { | ||
| yield return new MongoVectorSinkSettings(); | ||
| } | ||
| } |
1 change: 1 addition & 0 deletions
1
Extensions/Mongo/Cosmos.DataTransfer.MongoVectorExtension/Program.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| Console.WriteLine("Starting Mongo extension"); |
24 changes: 24 additions & 0 deletions
24
...Transfer.MongoVectorExtension/Properties/PublishProfiles/PublishToExtensionsFolder.pubxml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,24 @@ | ||
| <?xml version="1.0" encoding="utf-8"?> | ||
| <!-- | ||
| https://go.microsoft.com/fwlink/?LinkID=208121. | ||
| --> | ||
| <Project> | ||
| <PropertyGroup Condition=" '$(Configuration)' == 'Debug' "> | ||
| <Configuration>Debug</Configuration> | ||
| <Platform>Any CPU</Platform> | ||
| <PublishDir>..\..\..\Core\Cosmos.DataTransfer.Core\bin\Debug\net6.0\Extensions</PublishDir> | ||
| <PublishProtocol>FileSystem</PublishProtocol> | ||
| <_TargetId>Folder</_TargetId> | ||
| <TargetFramework>net6.0</TargetFramework> | ||
| <SelfContained>false</SelfContained> | ||
| </PropertyGroup> | ||
| <PropertyGroup Condition=" '$(Configuration)' != 'Debug' "> | ||
| <Configuration>Release</Configuration> | ||
| <Platform>Any CPU</Platform> | ||
| <PublishDir>..\..\..\Core\Cosmos.DataTransfer.Core\bin\Release\net6.0\Extensions</PublishDir> | ||
| <PublishProtocol>FileSystem</PublishProtocol> | ||
| <_TargetId>Folder</_TargetId> | ||
| <TargetFramework>net6.0</TargetFramework> | ||
| <SelfContained>false</SelfContained> | ||
| </PropertyGroup> | ||
| </Project> |
21 changes: 21 additions & 0 deletions
21
...nsions/Mongo/Cosmos.DataTransfer.MongoVectorExtension/Settings/MongoVectorSinkSettings.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| using System.ComponentModel.DataAnnotations; | ||
| using Cosmos.DataTransfer.MongoExtension.Settings; | ||
|
|
||
| namespace Cosmos.DataTransfer.MongoVectorExtension.Settings; | ||
| public class MongoVectorSinkSettings : MongoBaseSettings | ||
| { | ||
| [Required] | ||
| public string? Collection { get; set; } | ||
|
|
||
| public int? BatchSize { get; set; } | ||
|
|
||
| public bool? GenerateEmbedding { get; set; } | ||
|
|
||
| public string? OpenAIUrl { get; set; } | ||
| public string? OpenAIKey { get; set; } | ||
|
|
||
| // name of the deployment for text-embedding-ada-002 | ||
| public string? OpenAIDeploymentName { get; set; } | ||
| public string? SourcePropEmbedding { get; set; } | ||
| public string? DestPropEmbedding { get; set; } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
22 changes: 22 additions & 0 deletions
22
Extensions/PostgreSQL/Cosmos.DataTransfer.PostgresqlExtension.csproj
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,22 @@ | ||
| <Project Sdk="Microsoft.NET.Sdk"> | ||
|
|
||
| <PropertyGroup> | ||
| <OutputType>Exe</OutputType> | ||
| <TargetFramework>net6.0</TargetFramework> | ||
| <ImplicitUsings>enable</ImplicitUsings> | ||
| <Nullable>enable</Nullable> | ||
| </PropertyGroup> | ||
|
|
||
| <ItemGroup> | ||
| <PackageReference Include="Microsoft.Extensions.Configuration.Binder" Version="6.0.0" /> | ||
| <PackageReference Include="Npgsql" Version="7.0.6" /> | ||
| <PackageReference Include="System.ComponentModel.Composition" Version="7.0.0" /> | ||
| </ItemGroup> | ||
|
|
||
| <ItemGroup> | ||
| <ProjectReference Include="..\..\Interfaces\Cosmos.DataTransfer.Interfaces\Cosmos.DataTransfer.Interfaces.csproj" /> | ||
| </ItemGroup> | ||
| <Target Name="PublishToExtensionsFolder" AfterTargets="Build" Condition=" '$(Configuration)' == 'Debug' "> | ||
| <Exec Command="dotnet publish --configuration $(Configuration) --no-build -p:PublishProfile=FolderProfile" /> | ||
| </Target> | ||
| </Project> |
Oops, something went wrong.