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 all commits
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
224 changes: 149 additions & 75 deletions src/collections/generate/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,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 @@ -55,19 +57,19 @@ class GenerateManager<T, V> implements Generate<T, V> {
);
}

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

private async parseGroupByReply<C extends GenerativeConfigRuntime | undefined>(
opts: SearchOptions<T, V> | GroupByOptions<T> | undefined,
private async parseGroupByReply<RV, C extends GenerativeConfigRuntime | undefined>(
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, C>(reply);
? deserialize.generateGroupBy<T, RV>(reply)
: deserialize.generate<T, RV, C>(reply);
}

public fetchObjects<C extends GenerativeConfigRuntime | undefined = undefined>(
Expand All @@ -90,21 +92,29 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseReply(reply));
}

public bm25<C extends GenerativeConfigRuntime | undefined = undefined>(
public bm25<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
query: string,
generate: GenerateOptions<T, C>,
opts?: BaseBm25Options<T, V>
): Promise<GenerativeReturn<T, V, C>>;
public bm25<C extends GenerativeConfigRuntime | undefined = undefined>(
opts?: BaseBm25Options<T, I>
): Promise<GenerativeReturn<T, RV, C>>;
public bm25<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
query: string,
generate: GenerateOptions<T, C>,
opts: GroupByBm25Options<T, V>
): Promise<GenerativeGroupByReturn<T, V, C>>;
public bm25<C extends GenerativeConfigRuntime | undefined = undefined>(
query: string,
generate: GenerateOptions<T, C>,
opts?: Bm25Options<T, V>
): GenerateReturn<T, V, C> {
opts: GroupByBm25Options<T, I>
): Promise<GenerativeGroupByReturn<T, RV, C>>;
public bm25<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(query: string, generate: GenerateOptions<T, C>, opts?: Bm25Options<T, I>): GenerateReturn<T, RV, C> {
return Promise.all([
this.check.bm25(opts),
this.check.supportForSingleGroupedGenerative(),
Expand All @@ -121,21 +131,29 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

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

public nearImage<C extends GenerativeConfigRuntime | undefined = undefined>(
public nearImage<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
image: string | Buffer,
generate: GenerateOptions<T, C>,
opts?: BaseNearOptions<T, V>
): Promise<GenerativeReturn<T, V, C>>;
public nearImage<C extends GenerativeConfigRuntime | undefined = undefined>(
opts?: BaseNearOptions<T, V, I>
): Promise<GenerativeReturn<T, RV, C>>;
public nearImage<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
image: string | Buffer,
generate: GenerateOptions<T, C>,
opts: GroupByNearOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V, C>>;
public nearImage<C extends GenerativeConfigRuntime | undefined = undefined>(
opts: GroupByNearOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV, C>>;
public nearImage<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
image: string | Buffer,
generate: GenerateOptions<T, C>,
opts?: NearOptions<T, V>
): GenerateReturn<T, V, C> {
opts?: NearOptions<T, V, I>
): GenerateReturn<T, RV, C> {
return Promise.all([
this.check.nearSearch(opts),
this.check.supportForSingleGroupedGenerative(),
Expand All @@ -204,21 +234,29 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

public nearObject<C extends GenerativeConfigRuntime | undefined = undefined>(
id: string,
generate: GenerateOptions<T, C>,
opts?: BaseNearOptions<T, V>
): Promise<GenerativeReturn<T, V, C>>;
public nearObject<C extends GenerativeConfigRuntime | undefined = undefined>(
public nearObject<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
id: string,
generate: GenerateOptions<T, C>,
opts: GroupByNearOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V, C>>;
public nearObject<C extends GenerativeConfigRuntime | undefined = undefined>(
opts?: BaseNearOptions<T, V, I>
): Promise<GenerativeReturn<T, RV, C>>;
public nearObject<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
id: string,
generate: GenerateOptions<T, C>,
opts?: NearOptions<T, V>
): GenerateReturn<T, V, C> {
opts: GroupByNearOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV, C>>;
public nearObject<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(id: string, generate: GenerateOptions<T, C>, opts?: NearOptions<T, V, I>): GenerateReturn<T, RV, C> {
return Promise.all([
this.check.nearSearch(opts),
this.check.supportForSingleGroupedGenerative(),
Expand All @@ -242,21 +280,33 @@ class GenerateManager<T, V> implements Generate<T, V> {
.then((reply) => this.parseGroupByReply(opts, reply));
}

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

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

public nearMedia<C extends GenerativeConfigRuntime | undefined = undefined>(
public nearMedia<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
media: string | Buffer,
type: NearMediaType,
generate: GenerateOptions<T, C>,
opts?: BaseNearOptions<T, V>
): Promise<GenerativeReturn<T, V, C>>;
public nearMedia<C extends GenerativeConfigRuntime | undefined = undefined>(
opts?: BaseNearOptions<T, V, I>
): Promise<GenerativeReturn<T, RV, C>>;
public nearMedia<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
media: string | Buffer,
type: NearMediaType,
generate: GenerateOptions<T, C>,
opts: GroupByNearOptions<T, V>
): Promise<GenerativeGroupByReturn<T, V, C>>;
public nearMedia<C extends GenerativeConfigRuntime | undefined = undefined>(
opts: GroupByNearOptions<T, V, I>
): Promise<GenerativeGroupByReturn<T, RV, C>>;
public nearMedia<
I extends IncludeVector<V>,
RV extends ReturnVectors<V, I>,
C extends GenerativeConfigRuntime | undefined = undefined
>(
media: string | Buffer,
type: NearMediaType,
generate: GenerateOptions<T, C>,
opts?: NearOptions<T, V>
): GenerateReturn<T, V, C> {
opts?: NearOptions<T, V, I>
): GenerateReturn<T, RV, C> {
return Promise.all([
this.check.nearSearch(opts),
this.check.supportForSingleGroupedGenerative(),
Expand Down
Loading