feat(rust): openai embedding function #1275
Conversation
| .execute() | ||
| .await?; | ||
|
|
||
| // there is no equivalent to '.search(<query>)' yet |
There was a problem hiding this comment.
Is the plan to have an equivalent at some point in the future?
There was a problem hiding this comment.
IIRC, there were plans to add similar functionality to rust sdk
| // We can't use the FixedSizeListBuilder here because it always adds a null bitmap | ||
| // and we want to explicitly work with non-nullable arrays. |
There was a problem hiding this comment.
I see that OpenAi only supports non-nullable but will this builder hide the nulls by dropping the nullability (so we don't get an error when we should get an error)?
There was a problem hiding this comment.
no, we assert that the input array is non null, so we should never run into a scenario where we drop the nullability.
| _ => unreachable!("This should not happen. We already checked the data type."), | ||
| }; | ||
|
|
||
| let client = Client::with_config(creds); |
There was a problem hiding this comment.
Do we want to cache / pool this client at some point in the future? Although maybe that's more work than its worth given the amount of time it takes to calculate embeddings the connection establishment might be a small fraction anyways.
There was a problem hiding this comment.
that was my initial thought. We can definitely add this as an optimization later on if we need to.
| task::block_in_place(move || { | ||
| Handle::current().block_on(async { |
There was a problem hiding this comment.
Should compute_inner be an async function?
There was a problem hiding this comment.
I think we'd still need to spawn an async task for this regardless unless we made EmbeddingFunction::compute_source_embeddings async, but we can't do that without some significant refactoring.
We use the RecordBatchReader which isn't async. We'd likely have to refactor all instances of that to an async equivalent.
However, if you prefer the wrapping to happen outside of compute_inner, that's totally fine & i can make those changes.
There was a problem hiding this comment.
We don't have to worry about it right now. I am thinking we might eventually want compute_source_embeddings to be async though. The top-level add and query methods are already async so it seems odd to go async -> sync -> async.
Also, we might want to run embeddings in parallel at some point too. Although that might be best left to embedding function since different embeddings might have different preferences on batch size / parallel execution.
There was a problem hiding this comment.
We don't have to worry about it right now. I am thinking we might eventually want compute_source_embeddings to be async though. The top-level add and query methods are already async so it seems odd to go async -> sync -> async.
makes sense. Will update!
Also, we might want to run embeddings in parallel at some point too. Although that might be best left to embedding function since different embeddings might have different preferences on batch size / parallel execution.
I made a similar comment on an earlier PR. #1259 (comment)
github-actions
bot
added
enhancement
part of #994.
Adds the ability to use the openai embedding functions.
the example can be run by the following
which should output