Skip to content

Commit f7f525a

Browse files
tsmith023tsmith023
committed
Use type of includeVector to narrow vectors generic in returns
1 parent 47e8098 commit f7f525a

File tree

9 files changed

+499
-285
lines changed

9 files changed

+499
-285
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: 73 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,9 @@ import {
3131
GenerativeGroupByReturn,
3232
GenerativeReturn,
3333
GroupByOptions,
34+
ReturnVectors,
3435
} from '../types/index.js';
36+
import { IncludeVector } from '../types/internal.js';
3537
import { Generate } from './types.js';
3638

3739
class GenerateManager<T, V> implements Generate<T, V> {
@@ -53,19 +55,19 @@ class GenerateManager<T, V> implements Generate<T, V> {
5355
);
5456
}
5557

56-
private async parseReply(reply: SearchReply) {
58+
private async parseReply<RV>(reply: SearchReply) {
5759
const deserialize = await Deserialize.use(this.check.dbVersionSupport);
58-
return deserialize.generate<T, V>(reply);
60+
return deserialize.generate<T, RV>(reply);
5961
}
6062

61-
private async parseGroupByReply(
62-
opts: SearchOptions<T, V> | GroupByOptions<T> | undefined,
63+
private async parseGroupByReply<RV>(
64+
opts: SearchOptions<any, any> | GroupByOptions<any> | undefined,
6365
reply: SearchReply
6466
) {
6567
const deserialize = await Deserialize.use(this.check.dbVersionSupport);
6668
return Serialize.search.isGroupBy(opts)
67-
? deserialize.generateGroupBy<T, V>(reply)
68-
: deserialize.generate<T, V>(reply);
69+
? deserialize.generateGroupBy<T, RV>(reply)
70+
: deserialize.generate<T, RV>(reply);
6971
}
7072

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

88-
public bm25(
90+
public bm25<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
8991
query: string,
9092
generate: GenerateOptions<T>,
91-
opts?: BaseBm25Options<T, V>
92-
): Promise<GenerativeReturn<T, V>>;
93-
public bm25(
93+
opts?: BaseBm25Options<T, I>
94+
): Promise<GenerativeReturn<T, RV>>;
95+
public bm25<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
9496
query: string,
9597
generate: GenerateOptions<T>,
96-
opts: GroupByBm25Options<T, V>
97-
): Promise<GenerativeGroupByReturn<T, V>>;
98-
public bm25(query: string, generate: GenerateOptions<T>, opts?: Bm25Options<T, V>): GenerateReturn<T, V> {
98+
opts: GroupByBm25Options<T, I>
99+
): Promise<GenerativeGroupByReturn<T, RV>>;
100+
public bm25<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
101+
query: string,
102+
generate: GenerateOptions<T>,
103+
opts?: Bm25Options<T, I>
104+
): GenerateReturn<T, RV> {
99105
return this.check
100106
.bm25(opts)
101107
.then(({ search }) => ({
@@ -109,21 +115,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
109115
.then((reply) => this.parseGroupByReply(opts, reply));
110116
}
111117

112-
public hybrid(
118+
public hybrid<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
113119
query: string,
114120
generate: GenerateOptions<T>,
115-
opts?: BaseHybridOptions<T, V>
116-
): Promise<GenerativeReturn<T, V>>;
117-
public hybrid(
121+
opts?: BaseHybridOptions<T, V, I>
122+
): Promise<GenerativeReturn<T, RV>>;
123+
public hybrid<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
118124
query: string,
119125
generate: GenerateOptions<T>,
120-
opts: GroupByHybridOptions<T, V>
121-
): Promise<GenerativeGroupByReturn<T, V>>;
122-
public hybrid(
126+
opts: GroupByHybridOptions<T, V, I>
127+
): Promise<GenerativeGroupByReturn<T, RV>>;
128+
public hybrid<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
123129
query: string,
124130
generate: GenerateOptions<T>,
125-
opts?: HybridOptions<T, V>
126-
): GenerateReturn<T, V> {
131+
opts?: HybridOptions<T, V, I>
132+
): GenerateReturn<T, RV> {
127133
return this.check
128134
.hybridSearch(opts)
129135
.then(
@@ -154,21 +160,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
154160
.then((reply) => this.parseGroupByReply(opts, reply));
155161
}
156162

157-
public nearImage(
163+
public nearImage<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
158164
image: string | Buffer,
159165
generate: GenerateOptions<T>,
160-
opts?: BaseNearOptions<T, V>
161-
): Promise<GenerativeReturn<T, V>>;
162-
public nearImage(
166+
opts?: BaseNearOptions<T, V, I>
167+
): Promise<GenerativeReturn<T, RV>>;
168+
public nearImage<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
163169
image: string | Buffer,
164170
generate: GenerateOptions<T>,
165-
opts: GroupByNearOptions<T, V>
166-
): Promise<GenerativeGroupByReturn<T, V>>;
167-
public nearImage(
171+
opts: GroupByNearOptions<T, V, I>
172+
): Promise<GenerativeGroupByReturn<T, RV>>;
173+
public nearImage<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
168174
image: string | Buffer,
169175
generate: GenerateOptions<T>,
170-
opts?: NearOptions<T, V>
171-
): GenerateReturn<T, V> {
176+
opts?: NearOptions<T, V, I>
177+
): GenerateReturn<T, RV> {
172178
return this.check
173179
.nearSearch(opts)
174180
.then(async ({ search, supportsTargets, supportsWeightsForTargets }) => ({
@@ -189,21 +195,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
189195
.then((reply) => this.parseGroupByReply(opts, reply));
190196
}
191197

192-
public nearObject(
198+
public nearObject<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
193199
id: string,
194200
generate: GenerateOptions<T>,
195-
opts?: BaseNearOptions<T, V>
196-
): Promise<GenerativeReturn<T, V>>;
197-
public nearObject(
201+
opts?: BaseNearOptions<T, V, I>
202+
): Promise<GenerativeReturn<T, RV>>;
203+
public nearObject<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
198204
id: string,
199205
generate: GenerateOptions<T>,
200-
opts: GroupByNearOptions<T, V>
201-
): Promise<GenerativeGroupByReturn<T, V>>;
202-
public nearObject(
206+
opts: GroupByNearOptions<T, V, I>
207+
): Promise<GenerativeGroupByReturn<T, RV>>;
208+
public nearObject<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
203209
id: string,
204210
generate: GenerateOptions<T>,
205-
opts?: NearOptions<T, V>
206-
): GenerateReturn<T, V> {
211+
opts?: NearOptions<T, V, I>
212+
): GenerateReturn<T, RV> {
207213
return this.check
208214
.nearSearch(opts)
209215
.then(({ search, supportsTargets, supportsWeightsForTargets }) => ({
@@ -224,21 +230,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
224230
.then((reply) => this.parseGroupByReply(opts, reply));
225231
}
226232

227-
public nearText(
233+
public nearText<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
228234
query: string | string[],
229235
generate: GenerateOptions<T>,
230-
opts?: BaseNearTextOptions<T, V>
231-
): Promise<GenerativeReturn<T, V>>;
232-
public nearText(
236+
opts?: BaseNearTextOptions<T, V, I>
237+
): Promise<GenerativeReturn<T, RV>>;
238+
public nearText<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
233239
query: string | string[],
234240
generate: GenerateOptions<T>,
235-
opts: GroupByNearTextOptions<T, V>
236-
): Promise<GenerativeGroupByReturn<T, V>>;
237-
public nearText(
241+
opts: GroupByNearTextOptions<T, V, I>
242+
): Promise<GenerativeGroupByReturn<T, RV>>;
243+
public nearText<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
238244
query: string | string[],
239245
generate: GenerateOptions<T>,
240-
opts?: NearOptions<T, V>
241-
): GenerateReturn<T, V> {
246+
opts?: NearOptions<T, V, I>
247+
): GenerateReturn<T, RV> {
242248
return this.check
243249
.nearSearch(opts)
244250
.then(({ search, supportsTargets, supportsWeightsForTargets }) => ({
@@ -259,21 +265,21 @@ class GenerateManager<T, V> implements Generate<T, V> {
259265
.then((reply) => this.parseGroupByReply(opts, reply));
260266
}
261267

262-
public nearVector(
268+
public nearVector<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
263269
vector: number[],
264270
generate: GenerateOptions<T>,
265-
opts?: BaseNearOptions<T, V>
266-
): Promise<GenerativeReturn<T, V>>;
267-
public nearVector(
271+
opts?: BaseNearOptions<T, V, I>
272+
): Promise<GenerativeReturn<T, RV>>;
273+
public nearVector<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
268274
vector: number[],
269275
generate: GenerateOptions<T>,
270-
opts: GroupByNearOptions<T, V>
271-
): Promise<GenerativeGroupByReturn<T, V>>;
272-
public nearVector(
276+
opts: GroupByNearOptions<T, V, I>
277+
): Promise<GenerativeGroupByReturn<T, RV>>;
278+
public nearVector<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
273279
vector: number[],
274280
generate: GenerateOptions<T>,
275-
opts?: NearOptions<T, V>
276-
): GenerateReturn<T, V> {
281+
opts?: NearOptions<T, V, I>
282+
): GenerateReturn<T, RV> {
277283
return this.check
278284
.nearVector(vector, opts)
279285
.then(
@@ -304,24 +310,24 @@ class GenerateManager<T, V> implements Generate<T, V> {
304310
.then((reply) => this.parseGroupByReply(opts, reply));
305311
}
306312

307-
public nearMedia(
313+
public nearMedia<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
308314
media: string | Buffer,
309315
type: NearMediaType,
310316
generate: GenerateOptions<T>,
311-
opts?: BaseNearOptions<T, V>
312-
): Promise<GenerativeReturn<T, V>>;
313-
public nearMedia(
317+
opts?: BaseNearOptions<T, V, I>
318+
): Promise<GenerativeReturn<T, RV>>;
319+
public nearMedia<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
314320
media: string | Buffer,
315321
type: NearMediaType,
316322
generate: GenerateOptions<T>,
317-
opts: GroupByNearOptions<T, V>
318-
): Promise<GenerativeGroupByReturn<T, V>>;
319-
public nearMedia(
323+
opts: GroupByNearOptions<T, V, I>
324+
): Promise<GenerativeGroupByReturn<T, RV>>;
325+
public nearMedia<I extends IncludeVector<V>, RV extends ReturnVectors<V, I>>(
320326
media: string | Buffer,
321327
type: NearMediaType,
322328
generate: GenerateOptions<T>,
323-
opts?: NearOptions<T, V>
324-
): GenerateReturn<T, V> {
329+
opts?: NearOptions<T, V, I>
330+
): GenerateReturn<T, RV> {
325331
return this.check
326332
.nearSearch(opts)
327333
.then(({ search, supportsTargets, supportsWeightsForTargets }) => {

0 commit comments

Comments
 (0)