Skip to content

Commit

Permalink
Implement Vectorize GA binding changes
Browse files Browse the repository at this point in the history
  • Loading branch information
@ndisidore
ndisidore committed Jul 26, 2024
1 parent d7e4ea7 commit df5356e
Show file tree
Hide file tree
Showing 6 changed files with 176 additions and 118 deletions.
2 changes: 1 addition & 1 deletion src/cloudflare/internal/test/vectorize/vectorize-api-test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ export const test_vector_search_vector_query = {
returnMetadata: "indexed",
});
assert.equal(true, results.count > 0);
/** @type {VectorizeQueryMatches} */
/** @type {VectorizeMatches} */
const expected = {
matches: [
{
Expand Down
4 changes: 3 additions & 1 deletion src/cloudflare/internal/test/vectorize/vectorize-api-test.wd-test
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@ const unitTests :Workerd.Config = (
innerBindings = [(
name = "fetcher",
service = "vector-search-mock"
)],
),
(name = "indexId", text = "an-index"),
(name = "indexVersion", text = "v2")],
)
)
],
Expand Down
9 changes: 2 additions & 7 deletions src/cloudflare/internal/test/vectorize/vectorize-mock.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
// https://opensource.org/licenses/Apache-2.0

/** @type {Array<VectorizeQueryMatch>} */
/** @type {Array<VectorizeMatch>} */
const exampleVectorMatches = [
{
id: "b0daca4a-ffd8-4865-926b-e24800af2a2d",
Expand Down Expand Up @@ -97,13 +97,8 @@ export default {
) {
return Response.json({});
} else if (request.method === "POST" && pathname.endsWith("/query")) {
/** @type {VectorizeQueryOptions & {vector: number[], compat: { queryMetadataOptional: boolean }}} */
/** @type {VectorizeQueryOptions<VectorizeMetadataRetrievalLevel> & {vector: number[]}} */
const body = await request.json();
// check that the compatibility flags are set
if (!body.compat.queryMetadataOptional)
throw Error(
"expected to get `queryMetadataOptional` compat flag with a value of true"
);
let returnSet = structuredClone(exampleVectorMatches);
if (
body?.filter?.["text"] &&
Expand Down
241 changes: 134 additions & 107 deletions src/cloudflare/internal/vectorize-api.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,144 +16,166 @@ enum Operation {
VECTOR_DELETE = 5,
}

class VectorizeIndexImpl implements VectorizeIndex {
type VectorizeVersion = "v1" | "v2";

/*
* The Vectorize beta VectorizeIndex shares the same methods, so to keep things simple, they share one implementation.
* The types here are specific to Vectorize GA, but the types here don't actually matter as they are stripped away
* and not visible to end users.
*/
class VectorizeIndexImpl implements Vectorize {
public constructor(
private readonly fetcher: Fetcher,
private readonly indexId: string
private readonly indexId: string,
private readonly indexVersion: VectorizeVersion
) {}

public async describe(): Promise<VectorizeIndexDetails> {
const res = await this._send(
Operation.INDEX_GET,
`indexes/${this.indexId}`,
{
method: "GET",
}
);
public async describe(): Promise<VectorizeIndexInfo> {
const endpoint =
this.indexVersion === "v2" ? `info` : `binding/indexes/${this.indexId}`;
const res = await this._send(Operation.INDEX_GET, endpoint, {
method: "GET",
});

return await toJson<VectorizeIndexDetails>(res);
return await toJson<VectorizeIndexInfo>(res);
}

public async query(
vector: VectorFloatArray | number[],
options: VectorizeQueryOptions
options: VectorizeQueryOptions<VectorizeMetadataRetrievalLevel>
): Promise<VectorizeMatches> {
const compat = {
queryMetadataOptional: flags.vectorizeQueryMetadataOptional,
};
const res = await this._send(
Operation.VECTOR_QUERY,
`indexes/${this.indexId}/query`,
{
if (this.indexVersion === "v2") {
const res = await this._send(Operation.VECTOR_QUERY, `query`, {
method: "POST",
body: JSON.stringify({
...options,
vector: Array.isArray(vector) ? vector : Array.from(vector),
compat,
}),
headers: {
"content-type": "application/json",
accept: "application/json",
"cf-vector-search-query-compat": JSON.stringify(compat),
},
}
);

return await toJson<VectorizeMatches>(res);
});

return await toJson<VectorizeMatches>(res);
} else {
const compat = {
queryMetadataOptional: flags.vectorizeQueryMetadataOptional,
};
const res = await this._send(
Operation.VECTOR_QUERY,
`binding/indexes/${this.indexId}/query`,
{
method: "POST",
body: JSON.stringify({
...options,
vector: Array.isArray(vector) ? vector : Array.from(vector),
compat,
}),
headers: {
"content-type": "application/json",
accept: "application/json",
"cf-vector-search-query-compat": JSON.stringify(compat),
},
}
);

return await toJson<VectorizeMatches>(res);
}
}

public async insert(
vectors: VectorizeVector[]
): Promise<VectorizeVectorMutation> {
const res = await this._send(
Operation.VECTOR_INSERT,
`indexes/${this.indexId}/insert`,
{
method: "POST",
body: JSON.stringify({
vectors: vectors.map((vec) => ({
...vec,
values: Array.isArray(vec.values)
? vec.values
: Array.from(vec.values),
})),
}),
headers: {
"content-type": "application/json",
"cf-vector-search-dim-width": String(
vectors.length ? vectors[0]?.values?.length : 0
),
"cf-vector-search-dim-height": String(vectors.length),
accept: "application/json",
},
}
);

return await toJson<VectorizeVectorMutation>(res);
): Promise<VectorizeAsyncMutation> {
const endpoint =
this.indexVersion === "v2"
? `insert`
: `binding/indexes/${this.indexId}/insert`;
const res = await this._send(Operation.VECTOR_INSERT, endpoint, {
method: "POST",
body: JSON.stringify({
vectors: vectors.map((vec) => ({
...vec,
values: Array.isArray(vec.values)
? vec.values
: Array.from(vec.values),
})),
}),
headers: {
"content-type": "application/json",
"cf-vector-search-dim-width": String(
vectors.length ? vectors[0]?.values?.length : 0
),
"cf-vector-search-dim-height": String(vectors.length),
accept: "application/json",
},
});

return await toJson<VectorizeAsyncMutation>(res);
}

public async upsert(
vectors: VectorizeVector[]
): Promise<VectorizeVectorMutation> {
const res = await this._send(
Operation.VECTOR_UPSERT,
`indexes/${this.indexId}/upsert`,
{
method: "POST",
body: JSON.stringify({
vectors: vectors.map((vec) => ({
...vec,
values: Array.isArray(vec.values)
? vec.values
: Array.from(vec.values),
})),
}),
headers: {
"content-type": "application/json",
"cf-vector-search-dim-width": String(
vectors.length ? vectors[0]?.values?.length : 0
),
"cf-vector-search-dim-height": String(vectors.length),
accept: "application/json",
},
}
);

return await toJson<VectorizeVectorMutation>(res);
): Promise<VectorizeAsyncMutation> {
const endpoint =
this.indexVersion === "v2"
? `upsert`
: `binding/indexes/${this.indexId}/upsert`;
const res = await this._send(Operation.VECTOR_UPSERT, endpoint, {
method: "POST",
body: JSON.stringify({
vectors: vectors.map((vec) => ({
...vec,
values: Array.isArray(vec.values)
? vec.values
: Array.from(vec.values),
})),
}),
headers: {
"content-type": "application/json",
"cf-vector-search-dim-width": String(
vectors.length ? vectors[0]?.values?.length : 0
),
"cf-vector-search-dim-height": String(vectors.length),
accept: "application/json",
},
});

return await toJson<VectorizeAsyncMutation>(res);
}

public async getByIds(ids: string[]): Promise<VectorizeVector[]> {
const res = await this._send(
Operation.VECTOR_GET,
`indexes/${this.indexId}/getByIds`,
{
method: "POST",
body: JSON.stringify({ ids }),
headers: {
"content-type": "application/json",
accept: "application/json",
},
}
);
const endpoint =
this.indexVersion === "v2"
? `getByIds`
: `binding/indexes/${this.indexId}/getByIds`;
const res = await this._send(Operation.VECTOR_GET, endpoint, {
method: "POST",
body: JSON.stringify({ ids }),
headers: {
"content-type": "application/json",
accept: "application/json",
},
});

return await toJson<VectorizeVector[]>(res);
}

public async deleteByIds(ids: string[]): Promise<VectorizeVectorMutation> {
const res = await this._send(
Operation.VECTOR_DELETE,
`indexes/${this.indexId}/deleteByIds`,
{
method: "POST",
body: JSON.stringify({ ids }),
headers: {
"content-type": "application/json",
accept: "application/json",
},
}
);

return await toJson<VectorizeVectorMutation>(res);
public async deleteByIds(ids: string[]): Promise<VectorizeAsyncMutation> {
const endpoint =
this.indexVersion === "v2"
? `deleteByIds`
: `binding/indexes/${this.indexId}/deleteByIds`;
const res = await this._send(Operation.VECTOR_DELETE, endpoint, {
method: "POST",
body: JSON.stringify({ ids }),
headers: {
"content-type": "application/json",
accept: "application/json",
},
});

return await toJson<VectorizeAsyncMutation>(res);
}

private async _send(
Expand All @@ -162,7 +184,7 @@ class VectorizeIndexImpl implements VectorizeIndex {
init: RequestInit
): Promise<Response> {
const res = await this.fetcher.fetch(
`http://vector-search/binding/${endpoint}`, // `http://vector-search` is just a dummy host, the attached fetcher will receive the request
`http://vector-search/${endpoint}`, // `http://vector-search` is just a dummy host, the attached fetcher will receive the request
init
);
if (res.status !== 200) {
Expand Down Expand Up @@ -217,8 +239,13 @@ async function toJson<T = unknown>(response: Response): Promise<T> {
export function makeBinding(env: {
fetcher: Fetcher;
indexId: string;
}): VectorizeIndex {
return new VectorizeIndexImpl(env.fetcher, env.indexId);
indexVersion?: VectorizeVersion;
}): Vectorize {
return new VectorizeIndexImpl(
env.fetcher,
env.indexId,
env.indexVersion ?? "v1"
);
}

export default makeBinding;
19 changes: 18 additions & 1 deletion src/cloudflare/internal/vectorize.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,9 @@ type VectorizeIndexConfig =

/**
* Metadata about an existing index.
*
* This type is exclusively for the Vectorize **beta** and will be deprecated once Vectorize RC is released.
* See {@link VectorizeIndexInfo} for its post-beta equivalent.
*/
interface VectorizeIndexDetails {
/** The unique ID of the index */
Expand All @@ -105,6 +108,20 @@ interface VectorizeIndexDetails {
vectorsCount: number;
}

/**
* Metadata about an existing index.
*/
interface VectorizeIndexInfo {
/** The number of records containing vectors within the index. */
vectorsCount: number;
/** Number of dimensions the index has been configured for. */
dimensions: number;
/** ISO 8601 datetime of the last processed mutation on in the index. All changes before this mutation will be reflected in the index state. */
processedUpToDatetime: number;
/** UUIDv4 of the last mutation processed by the index. All changes before this mutation will be reflected in the index state. */
processedUpToMutation: number;
}

/**
* Represents a single vector value set along with its associated metadata.
*/
Expand Down Expand Up @@ -217,7 +234,7 @@ declare abstract class Vectorize {
* Get information about the currently bound index.
* @returns A promise that resolves with information about the current index.
*/
public describe(): Promise<VectorizeIndexDetails>;
public describe(): Promise<VectorizeIndexInfo>;
/**
* Use the provided vector to perform a similarity search across the index.
* @param vector Input vector that will be used to drive the similarity search.
Expand Down
Loading

0 comments on commit df5356e

Please sign in to comment.