forked from ravendb/ravendb
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
RavenDB-22076 Vector search client API
Showing
14 changed files
with
566 additions
and
55 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
139 changes: 139 additions & 0 deletions
139
src/Raven.Client/Documents/Queries/VectorFieldFactory.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,139 @@ | ||
using System; | ||
using System.Linq.Expressions; | ||
using Raven.Client.Documents.Conventions; | ||
using Raven.Client.Extensions; | ||
|
||
namespace Raven.Client.Documents.Queries; | ||
|
||
public interface IVectorFieldFactory<T> | ||
{ | ||
public IVectorEmbeddingTextField WithText(string fieldName); | ||
|
||
public IVectorEmbeddingTextField WithText(Expression<Func<T, string>> propertySelector); | ||
|
||
public IVectorEmbeddingField WithEmbedding(string fieldName, EmbeddingQuantizationType storedEmbeddingQuantization = EmbeddingQuantizationType.None); | ||
|
||
public IVectorEmbeddingField WithEmbedding(Expression<Func<T, string>> propertySelector, EmbeddingQuantizationType storedEmbeddingQuantization = EmbeddingQuantizationType.None); | ||
public IVectorEmbeddingField WithBase64(string fieldName, EmbeddingQuantizationType storedEmbeddingQuantization = EmbeddingQuantizationType.None); | ||
|
||
public IVectorEmbeddingField WithBase64(Expression<Func<T, string>> propertySelector, EmbeddingQuantizationType storedEmbeddingQuantization = EmbeddingQuantizationType.None); | ||
} | ||
|
||
public interface IVectorEmbeddingTextField | ||
{ | ||
|
||
} | ||
|
||
public interface IVectorEmbeddingField | ||
{ | ||
public IVectorEmbeddingField TargetQuantization(EmbeddingQuantizationType targetEmbeddingQuantization); | ||
} | ||
|
||
internal sealed class VectorEmbeddingFieldFactory<T> : IVectorFieldFactory<T>, IVectorEmbeddingTextField, IVectorEmbeddingField | ||
{ | ||
internal string FieldName { get; set; } | ||
internal EmbeddingQuantizationType SourceQuantizationType { get; set; } | ||
internal EmbeddingQuantizationType DestinationQuantizationType { get; set; } | ||
internal bool IsBase64Encoded { get; set; } | ||
|
||
IVectorEmbeddingTextField IVectorFieldFactory<T>.WithText(Expression<Func<T, string>> propertySelector) | ||
{ | ||
FieldName = propertySelector.ToPropertyPath(DocumentConventions.Default); | ||
|
||
return this; | ||
} | ||
|
||
IVectorEmbeddingTextField IVectorFieldFactory<T>.WithText(string fieldName) | ||
{ | ||
FieldName = fieldName; | ||
|
||
return this; | ||
} | ||
|
||
IVectorEmbeddingField IVectorFieldFactory<T>.WithEmbedding(string fieldName, EmbeddingQuantizationType storedEmbeddingQuantization) | ||
{ | ||
FieldName = fieldName; | ||
SourceQuantizationType = storedEmbeddingQuantization; | ||
|
||
return this; | ||
} | ||
|
||
IVectorEmbeddingField IVectorFieldFactory<T>.WithEmbedding(Expression<Func<T, string>> propertySelector, EmbeddingQuantizationType storedEmbeddingQuantization) | ||
{ | ||
FieldName = propertySelector.ToPropertyPath(DocumentConventions.Default); | ||
SourceQuantizationType = storedEmbeddingQuantization; | ||
|
||
return this; | ||
} | ||
|
||
IVectorEmbeddingField IVectorFieldFactory<T>.WithBase64(string fieldName, EmbeddingQuantizationType storedEmbeddingQuantization) | ||
{ | ||
FieldName = fieldName; | ||
SourceQuantizationType = storedEmbeddingQuantization; | ||
IsBase64Encoded = true; | ||
|
||
return this; | ||
} | ||
|
||
IVectorEmbeddingField IVectorFieldFactory<T>.WithBase64(Expression<Func<T, string>> propertySelector, EmbeddingQuantizationType storedEmbeddingQuantization) | ||
{ | ||
FieldName = propertySelector.ToPropertyPath(DocumentConventions.Default); | ||
SourceQuantizationType = storedEmbeddingQuantization; | ||
IsBase64Encoded = true; | ||
|
||
return this; | ||
} | ||
|
||
IVectorEmbeddingField IVectorEmbeddingField.TargetQuantization(EmbeddingQuantizationType targetEmbeddingQuantization) | ||
{ | ||
DestinationQuantizationType = targetEmbeddingQuantization; | ||
|
||
return this; | ||
} | ||
} | ||
|
||
////////////////////////// | ||
|
||
public interface IVectorEmbeddingTextFieldValueFactory | ||
{ | ||
public void ByText(string text); | ||
} | ||
|
||
public interface IVectorEmbeddingFieldValueFactory | ||
{ | ||
public void ByEmbedding(float[] embedding, EmbeddingQuantizationType queriedEmbeddingQuantization = EmbeddingQuantizationType.None); | ||
public void ByBase64(string base64Embedding, EmbeddingQuantizationType queriedEmbeddingQuantization = EmbeddingQuantizationType.None); | ||
} | ||
|
||
internal class VectorEmbeddingValueFactory : IVectorEmbeddingFieldValueFactory, IVectorEmbeddingTextFieldValueFactory | ||
{ | ||
public float[] Embedding { get; set; } | ||
public string Text { get; set; } | ||
public string Base64Embedding { get; set; } | ||
public EmbeddingQuantizationType EmbeddingQuantizationType { get; set; } | ||
|
||
void IVectorEmbeddingFieldValueFactory.ByEmbedding(float[] embedding, EmbeddingQuantizationType queriedEmbeddingQuantization) | ||
{ | ||
Embedding = embedding; | ||
EmbeddingQuantizationType = queriedEmbeddingQuantization; | ||
} | ||
|
||
void IVectorEmbeddingFieldValueFactory.ByBase64(string base64Embedding, EmbeddingQuantizationType queriedEmbeddingQuantization) | ||
{ | ||
Base64Embedding = base64Embedding; | ||
EmbeddingQuantizationType = queriedEmbeddingQuantization; | ||
} | ||
|
||
void IVectorEmbeddingTextFieldValueFactory.ByText(string text) | ||
{ | ||
Text = text; | ||
} | ||
} | ||
|
||
public enum EmbeddingQuantizationType | ||
{ | ||
None = 0, | ||
F32 = None, | ||
I8 = 1, | ||
I1 = 2 | ||
} |
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
90 changes: 90 additions & 0 deletions
90
src/Raven.Client/Documents/Session/Tokens/VectorSearchToken.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,90 @@ | ||
using System.Text; | ||
using Raven.Client.Documents.Queries; | ||
|
||
namespace Raven.Client.Documents.Session.Tokens; | ||
|
||
public sealed class VectorSearchToken : WhereToken | ||
{ | ||
private float SimilarityThreshold { get; set; } | ||
|
||
private EmbeddingQuantizationType SourceQuantizationType { get; set; } | ||
|
||
private EmbeddingQuantizationType TargetQuantizationType { get; set; } | ||
|
||
private EmbeddingQuantizationType VectorQuantizationType { get; set; } | ||
|
||
private bool IsSourceBase64Encoded { get; set; } | ||
|
||
private bool IsVectorBase64Encoded { get; set; } | ||
|
||
public VectorSearchToken(string fieldName, string parameterName, EmbeddingQuantizationType sourceQuantizationType, EmbeddingQuantizationType targetQuantizationType, EmbeddingQuantizationType vectorQuantizationType, bool isSourceBase64Encoded, bool isVectorBase64Encoded, float similarityThreshold) | ||
{ | ||
FieldName = fieldName; | ||
ParameterName = parameterName; | ||
|
||
SourceQuantizationType = sourceQuantizationType; | ||
TargetQuantizationType = targetQuantizationType; | ||
VectorQuantizationType = vectorQuantizationType; | ||
|
||
IsSourceBase64Encoded = isSourceBase64Encoded; | ||
IsVectorBase64Encoded = isVectorBase64Encoded; | ||
|
||
SimilarityThreshold = similarityThreshold; | ||
} | ||
|
||
public override void WriteTo(StringBuilder writer) | ||
{ | ||
bool explicitSourceQuantizationType = false; | ||
|
||
writer.Append("vector.search("); | ||
|
||
if (IsSourceBase64Encoded) | ||
writer.Append("base64("); | ||
|
||
if (SourceQuantizationType == EmbeddingQuantizationType.None) | ||
{ | ||
if (TargetQuantizationType != EmbeddingQuantizationType.None) | ||
{ | ||
// TODO | ||
writer.Append($"f32_{TargetQuantizationType.ToString().ToLower()}("); | ||
|
||
explicitSourceQuantizationType = true; | ||
} | ||
} | ||
|
||
else if (TargetQuantizationType != EmbeddingQuantizationType.None) | ||
{ | ||
writer.Append($"{SourceQuantizationType.ToString().ToLower()}("); | ||
|
||
explicitSourceQuantizationType = true; | ||
} | ||
|
||
writer.Append(FieldName); | ||
|
||
if (IsSourceBase64Encoded) | ||
writer.Append(')'); | ||
|
||
if (explicitSourceQuantizationType) | ||
writer.Append(')'); | ||
|
||
writer.Append(", "); | ||
|
||
if (IsVectorBase64Encoded) | ||
writer.Append("base64("); | ||
|
||
if (VectorQuantizationType != EmbeddingQuantizationType.F32) | ||
writer.Append($"{VectorQuantizationType.ToString().ToLower()}("); | ||
|
||
writer.Append($"${ParameterName}"); | ||
|
||
if (VectorQuantizationType != EmbeddingQuantizationType.F32) | ||
writer.Append(')'); | ||
|
||
if (IsVectorBase64Encoded) | ||
writer.Append(')'); | ||
|
||
writer.Append($", {SimilarityThreshold}"); | ||
|
||
writer.Append(')'); | ||
} | ||
} |
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
Oops, something went wrong.