Skip to content

Use type of includeVector to narrow vectors generic in returns #265

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 7, 2025
Merged

Use type of includeVector to narrow vectors generic in returns #265

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f7f525a
Use type of `includeVector` to narrow vectors generic in returns
tsmith023 Feb 14, 2025
a473598
Merge branch '1.29/colbert-multivector-support' of https://github.com…
tsmith023 Apr 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 12 additions & 8 deletions src/collections/collection/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ import { Iterator } from '../iterator/index.js';
import query, { Query } from '../query/index.js';
import sort, { Sort } from '../sort/index.js';
import tenants, { TenantBase, Tenants } from '../tenants/index.js';
import { QueryMetadata, QueryProperty, QueryReference } from '../types/index.js';
import { QueryMetadata, QueryProperty, QueryReference, ReturnVectors } from '../types/index.js';
import { IncludeVector } from '../types/internal.js';
import multiTargetVector, { MultiTargetVector } from '../vectors/multiTargetVector.js';

Expand Down Expand Up @@ -55,15 +55,19 @@ export interface Collection<T = undefined, N = string, V = undefined> {
* This iterator keeps a record of the last object that it returned to be used in each subsequent call to Weaviate.
* Once the collection is exhausted, the iterator exits.
*
* @typeParam I - The vector(s) to include in the response. If using named vectors, pass an array of strings to include only specific vectors.
* @typeParam RV - The vectors(s) to be returned in the response depending on the input in opts.includeVector.
* @param {IteratorOptions<T>} opts The options to use when fetching objects from Weaviate.
* @returns {Iterator<T>} An iterator over the objects in the collection as an async generator.
*
* @description If `return_properties` is not provided, all the properties of each object will be
* requested from Weaviate except for its vector as this is an expensive operation. Specify `include_vector`
* to request the vector back as well. In addition, if `return_references=None` then none of the references
* to request the vectors back as well. In addition, if `return_references=None` then none of the references
* are returned. Use `wvc.QueryReference` to specify which references to return.
*/
iterator: (opts?: IteratorOptions<T, V>) => Iterator<T, V>;
iterator: <I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
opts?: IteratorOptions<T, I>
) => Iterator<T, RV>;
/**
* Use this method to return the total number of objects in the collection.
*
Expand Down Expand Up @@ -95,8 +99,8 @@ export interface Collection<T = undefined, N = string, V = undefined> {
withTenant: <TT extends TenantBase>(tenant: string | TT) => Collection<T, N, V>;
}

export type IteratorOptions<T, V> = {
includeVector?: IncludeVector<V>;
export type IteratorOptions<T, I> = {
includeVector?: I;
returnMetadata?: QueryMetadata;
returnProperties?: QueryProperty<T>[];
returnReferences?: QueryReference<T>[];
Expand Down Expand Up @@ -146,10 +150,10 @@ const collection = <T, N, V>(
sort: sort<T>(),
tenants: tenants(connection, capitalizedName, dbVersionSupport),
exists: () => new ClassExists(connection).withClassName(capitalizedName).do(),
iterator: (opts?: IteratorOptions<T, V>) =>
new Iterator<T, V>((limit: number, after?: string) =>
iterator: <I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(opts?: IteratorOptions<T, I>) =>
new Iterator<T, RV>((limit: number, after?: string) =>
queryCollection
.fetchObjects({
.fetchObjects<I, RV>({
limit,
after,
includeVector: opts?.includeVector,
Expand Down
140 changes: 73 additions & 67 deletions src/collections/generate/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@ import {
GenerativeGroupByReturn,
GenerativeReturn,
GroupByOptions,
ReturnVectors,
} from '../types/index.js';
import { IncludeVector } from '../types/internal.js';
import { Generate } from './types.js';

class GenerateManager<T, V> implements Generate<T, V> {
Expand All @@ -53,19 +55,19 @@ class GenerateManager<T, V> implements Generate<T, V> {
);
}

private async parseReply(reply: SearchReply) {
private async parseReply<RV>(reply: SearchReply) {
const deserialize = await Deserialize.use(this.check.dbVersionSupport);
return deserialize.generate<T, V>(reply);
return deserialize.generate<T, RV>(reply);
}

private async parseGroupByReply(
opts: SearchOptions<T, V> | GroupByOptions<T> | undefined,
private async parseGroupByReply<RV>(
opts: SearchOptions<any, any> | GroupByOptions<any> | undefined,
reply: SearchReply
) {
const deserialize = await Deserialize.use(this.check.dbVersionSupport);
return Serialize.search.isGroupBy(opts)
? deserialize.generateGroupBy<T, V>(reply)
: deserialize.generate<T, V>(reply);
? deserialize.generateGroupBy<T, RV>(reply)
: deserialize.generate<T, RV>(reply);
}

public fetchObjects(
Expand All @@ -85,17 +87,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseReply(reply));
}

public bm25(
public bm25<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string,
generate: GenerateOptions<T>,
opts?: BaseBm25Options<T, V>
): Promise<GenerativeReturn<T, V>>;
public bm25(
opts?: BaseBm25Options<T, I>
): Promise<GenerativeReturn<T, RV>>;
public bm25<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string,
generate: GenerateOptions<T>,
opts: GroupByBm25Options<T, V>
): Promise<GenerativeGroupByReturn<T, V>>;
public bm25(query: string, generate: GenerateOptions<T>, opts?: Bm25Options<T, V>): GenerateReturn<T, V> {
opts: GroupByBm25Options<T, I>
): Promise<GenerativeGroupByReturn<T, RV>>;
public bm25<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string,
generate: GenerateOptions<T>,
opts?: Bm25Options<T, I>
): GenerateReturn<T, RV> {
return this.check
.bm25(opts)
.then(({ search }) => ({
Expand All @@ -109,21 +115,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

public hybrid(
public hybrid<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string,
generate: GenerateOptions<T>,
opts?: BaseHybridOptions<T, V>
): Promise<GenerativeReturn<T, V>>;
public hybrid(
opts?: BaseHybridOptions<T, V, I>
): Promise<GenerativeReturn<T, RV>>;
public hybrid<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string,
generate: GenerateOptions<T>,
opts: GroupByHybridOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V>>;
public hybrid(
opts: GroupByHybridOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV>>;
public hybrid<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string,
generate: GenerateOptions<T>,
opts?: HybridOptions<T, V>
): GenerateReturn<T, V> {
opts?: HybridOptions<T, V, I>
): GenerateReturn<T, RV> {
return this.check
.hybridSearch(opts)
.then(
Expand Down Expand Up @@ -154,21 +160,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

public nearImage(
public nearImage<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
image: string | Buffer,
generate: GenerateOptions<T>,
opts?: BaseNearOptions<T, V>
): Promise<GenerativeReturn<T, V>>;
public nearImage(
opts?: BaseNearOptions<T, V, I>
): Promise<GenerativeReturn<T, RV>>;
public nearImage<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
image: string | Buffer,
generate: GenerateOptions<T>,
opts: GroupByNearOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V>>;
public nearImage(
opts: GroupByNearOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV>>;
public nearImage<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
image: string | Buffer,
generate: GenerateOptions<T>,
opts?: NearOptions<T, V>
): GenerateReturn<T, V> {
opts?: NearOptions<T, V, I>
): GenerateReturn<T, RV> {
return this.check
.nearSearch(opts)
.then(async ({ search, supportsTargets, supportsWeightsForTargets }) => ({
Expand All @@ -189,21 +195,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

public nearObject(
public nearObject<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
id: string,
generate: GenerateOptions<T>,
opts?: BaseNearOptions<T, V>
): Promise<GenerativeReturn<T, V>>;
public nearObject(
opts?: BaseNearOptions<T, V, I>
): Promise<GenerativeReturn<T, RV>>;
public nearObject<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
id: string,
generate: GenerateOptions<T>,
opts: GroupByNearOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V>>;
public nearObject(
opts: GroupByNearOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV>>;
public nearObject<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
id: string,
generate: GenerateOptions<T>,
opts?: NearOptions<T, V>
): GenerateReturn<T, V> {
opts?: NearOptions<T, V, I>
): GenerateReturn<T, RV> {
return this.check
.nearSearch(opts)
.then(({ search, supportsTargets, supportsWeightsForTargets }) => ({
Expand All @@ -224,21 +230,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

public nearText(
public nearText<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string | string[],
generate: GenerateOptions<T>,
opts?: BaseNearTextOptions<T, V>
): Promise<GenerativeReturn<T, V>>;
public nearText(
opts?: BaseNearTextOptions<T, V, I>
): Promise<GenerativeReturn<T, RV>>;
public nearText<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string | string[],
generate: GenerateOptions<T>,
opts: GroupByNearTextOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V>>;
public nearText(
opts: GroupByNearTextOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV>>;
public nearText<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
query: string | string[],
generate: GenerateOptions<T>,
opts?: NearOptions<T, V>
): GenerateReturn<T, V> {
opts?: NearOptions<T, V, I>
): GenerateReturn<T, RV> {
return this.check
.nearSearch(opts)
.then(({ search, supportsTargets, supportsWeightsForTargets }) => ({
Expand All @@ -259,21 +265,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

public nearVector(
public nearVector<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
vector: number[],
generate: GenerateOptions<T>,
opts?: BaseNearOptions<T, V>
): Promise<GenerativeReturn<T, V>>;
public nearVector(
opts?: BaseNearOptions<T, V, I>
): Promise<GenerativeReturn<T, RV>>;
public nearVector<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
vector: number[],
generate: GenerateOptions<T>,
opts: GroupByNearOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V>>;
public nearVector(
opts: GroupByNearOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV>>;
public nearVector<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
vector: number[],
generate: GenerateOptions<T>,
opts?: NearOptions<T, V>
): GenerateReturn<T, V> {
opts?: NearOptions<T, V, I>
): GenerateReturn<T, RV> {
return this.check
.nearVector(vector, opts)
.then(
Expand Down Expand Up @@ -304,24 +310,24 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

public nearMedia(
public nearMedia<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
media: string | Buffer,
type: NearMediaType,
generate: GenerateOptions<T>,
opts?: BaseNearOptions<T, V>
): Promise<GenerativeReturn<T, V>>;
public nearMedia(
opts?: BaseNearOptions<T, V, I>
): Promise<GenerativeReturn<T, RV>>;
public nearMedia<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
media: string | Buffer,
type: NearMediaType,
generate: GenerateOptions<T>,
opts: GroupByNearOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V>>;
public nearMedia(
opts: GroupByNearOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV>>;
public nearMedia<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
media: string | Buffer,
type: NearMediaType,
generate: GenerateOptions<T>,
opts?: NearOptions<T, V>
): GenerateReturn<T, V> {
opts?: NearOptions<T, V, I>
): GenerateReturn<T, RV> {
return this.check
.nearSearch(opts)
.then(({ search, supportsTargets, supportsWeightsForTargets }) => {
Expand Down
Loading