Skip to content

Commit 382a407

Browse files
tsmith023tsmith023
authored
Use type of includeVector to narrow vectors generic in returns (#265)
1 parent 2eef48b commit 382a407

File tree

9 files changed

+641
-299
lines changed

9 files changed

+641
-299
lines changed

‎src/collections/collection/index.ts‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import { Iterator } from '../iterator/index.js';
1414
import query, { Query } from '../query/index.js';
1515
import sort, { Sort } from '../sort/index.js';
1616
import tenants, { TenantBase, Tenants } from '../tenants/index.js';
17-
import { QueryMetadata, QueryProperty, QueryReference } from '../types/index.js';
17+
import { QueryMetadata, QueryProperty, QueryReference, ReturnVectors } from '../types/index.js';
1818
import { IncludeVector } from '../types/internal.js';
1919
import multiTargetVector, { MultiTargetVector } from '../vectors/multiTargetVector.js';
2020

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

98-
export type IteratorOptions<T, V> = {
99-
includeVector?: IncludeVector<V>;
102+
export type IteratorOptions<T, I> = {
103+
includeVector?: I;
100104
returnMetadata?: QueryMetadata;
101105
returnProperties?: QueryProperty<T>[];
102106
returnReferences?: QueryReference<T>[];
@@ -146,10 +150,10 @@ const collection = <T, N, V>(
146150
sort: sort<T>(),
147151
tenants: tenants(connection, capitalizedName, dbVersionSupport),
148152
exists: () => new ClassExists(connection).withClassName(capitalizedName).do(),
149-
iterator: (opts?: IteratorOptions<T, V>) =>
150-
new Iterator<T, V>((limit: number, after?: string) =>
153+
iterator: <I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(opts?: IteratorOptions<T, I>) =>
154+
new Iterator<T, RV>((limit: number, after?: string) =>
151155
queryCollection
152-
.fetchObjects({
156+
.fetchObjects<I, RV>({
153157
limit,
154158
after,
155159
includeVector: opts?.includeVector,

‎src/collections/generate/index.ts‎

Lines changed: 149 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,9 @@ import {
3333
GenerativeGroupByReturn,
3434
GenerativeReturn,
3535
GroupByOptions,
36+
ReturnVectors,
3637
} from '../types/index.js';
38+
import { IncludeVector } from '../types/internal.js';
3739
import { Generate } from './types.js';
3840

3941
class GenerateManager<T, V> implements Generate<T, V> {
@@ -55,19 +57,19 @@ class GenerateManager<T, V> implements Generate<T, V> {
5557
);
5658
}
5759

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

63-
private async parseGroupByReply<C extends GenerativeConfigRuntime | undefined>(
64-
opts: SearchOptions<T, V> | GroupByOptions<T> | undefined,
65+
private async parseGroupByReply<RV, C extends GenerativeConfigRuntime | undefined>(
66+
opts: SearchOptions<any, any> | GroupByOptions<any> | undefined,
6567
reply: SearchReply
6668
) {
6769
const deserialize = await Deserialize.use(this.check.dbVersionSupport);
6870
return Serialize.search.isGroupBy(opts)
69-
? deserialize.generateGroupBy<T, V>(reply)
70-
: deserialize.generate<T, V, C>(reply);
71+
? deserialize.generateGroupBy<T, RV>(reply)
72+
: deserialize.generate<T, RV, C>(reply);
7173
}
7274

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

93-
public bm25<C extends GenerativeConfigRuntime | undefined = undefined>(
95+
public bm25<
96+
I extends IncludeVector<V>,
97+
RV extends ReturnVectors<V, I>,
98+
C extends GenerativeConfigRuntime | undefined = undefined
99+
>(
94100
query: string,
95101
generate: GenerateOptions<T, C>,
96-
opts?: BaseBm25Options<T, V>
97-
): Promise<GenerativeReturn<T, V, C>>;
98-
public bm25<C extends GenerativeConfigRuntime | undefined = undefined>(
102+
opts?: BaseBm25Options<T, I>
103+
): Promise<GenerativeReturn<T, RV, C>>;
104+
public bm25<
105+
I extends IncludeVector<V>,
106+
RV extends ReturnVectors<V, I>,
107+
C extends GenerativeConfigRuntime | undefined = undefined
108+
>(
99109
query: string,
100110
generate: GenerateOptions<T, C>,
101-
opts: GroupByBm25Options<T, V>
102-
): Promise<GenerativeGroupByReturn<T, V, C>>;
103-
public bm25<C extends GenerativeConfigRuntime | undefined = undefined>(
104-
query: string,
105-
generate: GenerateOptions<T, C>,
106-
opts?: Bm25Options<T, V>
107-
): GenerateReturn<T, V, C> {
111+
opts: GroupByBm25Options<T, I>
112+
): Promise<GenerativeGroupByReturn<T, RV, C>>;
113+
public bm25<
114+
I extends IncludeVector<V>,
115+
RV extends ReturnVectors<V, I>,
116+
C extends GenerativeConfigRuntime | undefined = undefined
117+
>(query: string, generate: GenerateOptions<T, C>, opts?: Bm25Options<T, I>): GenerateReturn<T, RV, C> {
108118
return Promise.all([
109119
this.check.bm25(opts),
110120
this.check.supportForSingleGroupedGenerative(),
@@ -121,21 +131,29 @@ class GenerateManager<T, V> implements Generate<T, V> {
121131
.then((reply) => this.parseGroupByReply(opts, reply));
122132
}
123133

124-
public hybrid<C extends GenerativeConfigRuntime | undefined = undefined>(
125-
query: string,
126-
generate: GenerateOptions<T, C>,
127-
opts?: BaseHybridOptions<T, V>
128-
): Promise<GenerativeReturn<T, V, C>>;
129-
public hybrid<C extends GenerativeConfigRuntime | undefined = undefined>(
134+
public hybrid<
135+
I extends IncludeVector<V>,
136+
RV extends ReturnVectors<V, I>,
137+
C extends GenerativeConfigRuntime | undefined = undefined
138+
>(
130139
query: string,
131140
generate: GenerateOptions<T, C>,
132-
opts: GroupByHybridOptions<T, V>
133-
): Promise<GenerativeGroupByReturn<T, V, C>>;
134-
public hybrid<C extends GenerativeConfigRuntime | undefined = undefined>(
141+
opts?: BaseHybridOptions<T, V, I>
142+
): Promise<GenerativeReturn<T, RV, C>>;
143+
public hybrid<
144+
I extends IncludeVector<V>,
145+
RV extends ReturnVectors<V, I>,
146+
C extends GenerativeConfigRuntime | undefined = undefined
147+
>(
135148
query: string,
136149
generate: GenerateOptions<T, C>,
137-
opts?: HybridOptions<T, V>
138-
): GenerateReturn<T, V, C> {
150+
opts: GroupByHybridOptions<T, V, I>
151+
): Promise<GenerativeGroupByReturn<T, RV, C>>;
152+
public hybrid<
153+
I extends IncludeVector<V>,
154+
RV extends ReturnVectors<V, I>,
155+
C extends GenerativeConfigRuntime | undefined = undefined
156+
>(query: string, generate: GenerateOptions<T, C>, opts?: HybridOptions<T, V, I>): GenerateReturn<T, RV, C> {
139157
return Promise.all([
140158
this.check.hybridSearch(opts),
141159
this.check.supportForSingleGroupedGenerative(),
@@ -166,21 +184,33 @@ class GenerateManager<T, V> implements Generate<T, V> {
166184
.then((reply) => this.parseGroupByReply(opts, reply));
167185
}
168186

169-
public nearImage<C extends GenerativeConfigRuntime | undefined = undefined>(
187+
public nearImage<
188+
I extends IncludeVector<V>,
189+
RV extends ReturnVectors<V, I>,
190+
C extends GenerativeConfigRuntime | undefined = undefined
191+
>(
170192
image: string | Buffer,
171193
generate: GenerateOptions<T, C>,
172-
opts?: BaseNearOptions<T, V>
173-
): Promise<GenerativeReturn<T, V, C>>;
174-
public nearImage<C extends GenerativeConfigRuntime | undefined = undefined>(
194+
opts?: BaseNearOptions<T, V, I>
195+
): Promise<GenerativeReturn<T, RV, C>>;
196+
public nearImage<
197+
I extends IncludeVector<V>,
198+
RV extends ReturnVectors<V, I>,
199+
C extends GenerativeConfigRuntime | undefined = undefined
200+
>(
175201
image: string | Buffer,
176202
generate: GenerateOptions<T, C>,
177-
opts: GroupByNearOptions<T, V>
178-
): Promise<GenerativeGroupByReturn<T, V, C>>;
179-
public nearImage<C extends GenerativeConfigRuntime | undefined = undefined>(
203+
opts: GroupByNearOptions<T, V, I>
204+
): Promise<GenerativeGroupByReturn<T, RV, C>>;
205+
public nearImage<
206+
I extends IncludeVector<V>,
207+
RV extends ReturnVectors<V, I>,
208+
C extends GenerativeConfigRuntime | undefined = undefined
209+
>(
180210
image: string | Buffer,
181211
generate: GenerateOptions<T, C>,
182-
opts?: NearOptions<T, V>
183-
): GenerateReturn<T, V, C> {
212+
opts?: NearOptions<T, V, I>
213+
): GenerateReturn<T, RV, C> {
184214
return Promise.all([
185215
this.check.nearSearch(opts),
186216
this.check.supportForSingleGroupedGenerative(),
@@ -204,21 +234,29 @@ class GenerateManager<T, V> implements Generate<T, V> {
204234
.then((reply) => this.parseGroupByReply(opts, reply));
205235
}
206236

207-
public nearObject<C extends GenerativeConfigRuntime | undefined = undefined>(
208-
id: string,
209-
generate: GenerateOptions<T, C>,
210-
opts?: BaseNearOptions<T, V>
211-
): Promise<GenerativeReturn<T, V, C>>;
212-
public nearObject<C extends GenerativeConfigRuntime | undefined = undefined>(
237+
public nearObject<
238+
I extends IncludeVector<V>,
239+
RV extends ReturnVectors<V, I>,
240+
C extends GenerativeConfigRuntime | undefined = undefined
241+
>(
213242
id: string,
214243
generate: GenerateOptions<T, C>,
215-
opts: GroupByNearOptions<T, V>
216-
): Promise<GenerativeGroupByReturn<T, V, C>>;
217-
public nearObject<C extends GenerativeConfigRuntime | undefined = undefined>(
244+
opts?: BaseNearOptions<T, V, I>
245+
): Promise<GenerativeReturn<T, RV, C>>;
246+
public nearObject<
247+
I extends IncludeVector<V>,
248+
RV extends ReturnVectors<V, I>,
249+
C extends GenerativeConfigRuntime | undefined = undefined
250+
>(
218251
id: string,
219252
generate: GenerateOptions<T, C>,
220-
opts?: NearOptions<T, V>
221-
): GenerateReturn<T, V, C> {
253+
opts: GroupByNearOptions<T, V, I>
254+
): Promise<GenerativeGroupByReturn<T, RV, C>>;
255+
public nearObject<
256+
I extends IncludeVector<V>,
257+
RV extends ReturnVectors<V, I>,
258+
C extends GenerativeConfigRuntime | undefined = undefined
259+
>(id: string, generate: GenerateOptions<T, C>, opts?: NearOptions<T, V, I>): GenerateReturn<T, RV, C> {
222260
return Promise.all([
223261
this.check.nearSearch(opts),
224262
this.check.supportForSingleGroupedGenerative(),
@@ -242,21 +280,33 @@ class GenerateManager<T, V> implements Generate<T, V> {
242280
.then((reply) => this.parseGroupByReply(opts, reply));
243281
}
244282

245-
public nearText<C extends GenerativeConfigRuntime | undefined = undefined>(
283+
public nearText<
284+
I extends IncludeVector<V>,
285+
RV extends ReturnVectors<V, I>,
286+
C extends GenerativeConfigRuntime | undefined = undefined
287+
>(
246288
query: string | string[],
247289
generate: GenerateOptions<T, C>,
248-
opts?: BaseNearTextOptions<T, V>
249-
): Promise<GenerativeReturn<T, V, C>>;
250-
public nearText<C extends GenerativeConfigRuntime | undefined = undefined>(
290+
opts?: BaseNearTextOptions<T, V, I>
291+
): Promise<GenerativeReturn<T, RV, C>>;
292+
public nearText<
293+
I extends IncludeVector<V>,
294+
RV extends ReturnVectors<V, I>,
295+
C extends GenerativeConfigRuntime | undefined = undefined
296+
>(
251297
query: string | string[],
252298
generate: GenerateOptions<T, C>,
253-
opts: GroupByNearTextOptions<T, V>
254-
): Promise<GenerativeGroupByReturn<T, V, C>>;
255-
public nearText<C extends GenerativeConfigRuntime | undefined = undefined>(
299+
opts: GroupByNearTextOptions<T, V, I>
300+
): Promise<GenerativeGroupByReturn<T, RV, C>>;
301+
public nearText<
302+
I extends IncludeVector<V>,
303+
RV extends ReturnVectors<V, I>,
304+
C extends GenerativeConfigRuntime | undefined = undefined
305+
>(
256306
query: string | string[],
257307
generate: GenerateOptions<T, C>,
258-
opts?: NearOptions<T, V>
259-
): GenerateReturn<T, V, C> {
308+
opts?: NearOptions<T, V, I>
309+
): GenerateReturn<T, RV, C> {
260310
return Promise.all([
261311
this.check.nearSearch(opts),
262312
this.check.supportForSingleGroupedGenerative(),
@@ -280,21 +330,33 @@ class GenerateManager<T, V> implements Generate<T, V> {
280330
.then((reply) => this.parseGroupByReply(opts, reply));
281331
}
282332

283-
public nearVector<C extends GenerativeConfigRuntime | undefined = undefined>(
333+
public nearVector<
334+
I extends IncludeVector<V>,
335+
RV extends ReturnVectors<V, I>,
336+
C extends GenerativeConfigRuntime | undefined = undefined
337+
>(
284338
vector: number[],
285339
generate: GenerateOptions<T, C>,
286-
opts?: BaseNearOptions<T, V>
287-
): Promise<GenerativeReturn<T, V, C>>;
288-
public nearVector<C extends GenerativeConfigRuntime | undefined = undefined>(
340+
opts?: BaseNearOptions<T, V, I>
341+
): Promise<GenerativeReturn<T, RV, C>>;
342+
public nearVector<
343+
I extends IncludeVector<V>,
344+
RV extends ReturnVectors<V, I>,
345+
C extends GenerativeConfigRuntime | undefined = undefined
346+
>(
289347
vector: number[],
290348
generate: GenerateOptions<T, C>,
291-
opts: GroupByNearOptions<T, V>
292-
): Promise<GenerativeGroupByReturn<T, V, C>>;
293-
public nearVector<C extends GenerativeConfigRuntime | undefined = undefined>(
349+
opts: GroupByNearOptions<T, V, I>
350+
): Promise<GenerativeGroupByReturn<T, RV, C>>;
351+
public nearVector<
352+
I extends IncludeVector<V>,
353+
RV extends ReturnVectors<V, I>,
354+
C extends GenerativeConfigRuntime | undefined = undefined
355+
>(
294356
vector: number[],
295357
generate: GenerateOptions<T, C>,
296-
opts?: NearOptions<T, V>
297-
): GenerateReturn<T, V, C> {
358+
opts?: NearOptions<T, V, I>
359+
): GenerateReturn<T, RV, C> {
298360
return Promise.all([
299361
this.check.nearVector(vector, opts),
300362
this.check.supportForSingleGroupedGenerative(),
@@ -325,24 +387,36 @@ class GenerateManager<T, V> implements Generate<T, V> {
325387
.then((reply) => this.parseGroupByReply(opts, reply));
326388
}
327389

328-
public nearMedia<C extends GenerativeConfigRuntime | undefined = undefined>(
390+
public nearMedia<
391+
I extends IncludeVector<V>,
392+
RV extends ReturnVectors<V, I>,
393+
C extends GenerativeConfigRuntime | undefined = undefined
394+
>(
329395
media: string | Buffer,
330396
type: NearMediaType,
331397
generate: GenerateOptions<T, C>,
332-
opts?: BaseNearOptions<T, V>
333-
): Promise<GenerativeReturn<T, V, C>>;
334-
public nearMedia<C extends GenerativeConfigRuntime | undefined = undefined>(
398+
opts?: BaseNearOptions<T, V, I>
399+
): Promise<GenerativeReturn<T, RV, C>>;
400+
public nearMedia<
401+
I extends IncludeVector<V>,
402+
RV extends ReturnVectors<V, I>,
403+
C extends GenerativeConfigRuntime | undefined = undefined
404+
>(
335405
media: string | Buffer,
336406
type: NearMediaType,
337407
generate: GenerateOptions<T, C>,
338-
opts: GroupByNearOptions<T, V>
339-
): Promise<GenerativeGroupByReturn<T, V, C>>;
340-
public nearMedia<C extends GenerativeConfigRuntime | undefined = undefined>(
408+
opts: GroupByNearOptions<T, V, I>
409+
): Promise<GenerativeGroupByReturn<T, RV, C>>;
410+
public nearMedia<
411+
I extends IncludeVector<V>,
412+
RV extends ReturnVectors<V, I>,
413+
C extends GenerativeConfigRuntime | undefined = undefined
414+
>(
341415
media: string | Buffer,
342416
type: NearMediaType,
343417
generate: GenerateOptions<T, C>,
344-
opts?: NearOptions<T, V>
345-
): GenerateReturn<T, V, C> {
418+
opts?: NearOptions<T, V, I>
419+
): GenerateReturn<T, RV, C> {
346420
return Promise.all([
347421
this.check.nearSearch(opts),
348422
this.check.supportForSingleGroupedGenerative(),

0 commit comments

Comments
 (0)