Skip to content

Commit b546836

Browse files
committed
Use more strict type guards in ai package
Turned on type-checked linting and fixed the low-hanging fruit of issues that arose. We can't yet add these lint rules in more places, but I wanted to at least get the benefit of these refactors before abandoning them for a while until we have time to address the issues.
1 parent 6420289 commit b546836

File tree

1 file changed

+44
-18
lines changed

1 file changed

+44
-18
lines changed

packages/ai/src/core.ts

+44-18
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import type { Client } from "gel";
22
import { EventSourceParserStream } from "eventsource-parser/stream";
33

4-
import type { ResolvedConnectConfig } from "gel/dist/conUtils.js";
54
import {
65
getAuthenticatedFetch,
76
type AuthenticatedFetch,
@@ -46,9 +45,7 @@ export class RAGClient {
4645
}
4746

4847
private static async getAuthenticatedFetch(client: Client) {
49-
const connectConfig: ResolvedConnectConfig = (
50-
await (client as any).pool._getNormalizedConnectConfig()
51-
).connectionParams;
48+
const connectConfig = await client.resolveConnectionParams();
5249

5350
return getAuthenticatedFetch(connectConfig, httpSCRAMAuth, "ext/ai/");
5451
}
@@ -110,7 +107,7 @@ export class RAGClient {
110107
!providedPrompt && {
111108
name: "builtin::rag-default",
112109
}),
113-
custom: [...(this.options.prompt?.custom || []), ...messages],
110+
custom: [...(this.options.prompt?.custom ?? []), ...messages],
114111
},
115112
query: [...messages].reverse().find((msg) => msg.role === "user")!
116113
.content[0].text,
@@ -139,18 +136,9 @@ export class RAGClient {
139136
);
140137
}
141138

142-
const data = await res.json();
139+
const data: unknown = await res.json();
143140

144-
if (
145-
!data ||
146-
typeof data !== "object" ||
147-
typeof data.response !== "string"
148-
) {
149-
throw new Error(
150-
"Expected response to be an object with response key of type string",
151-
);
152-
}
153-
return data.response;
141+
return parseResponse(data);
154142
}
155143

156144
streamRag(
@@ -190,6 +178,7 @@ export class RAGClient {
190178
}
191179
},
192180
then<TResult1 = Response, TResult2 = never>(
181+
/* eslint-disable @typescript-eslint/no-duplicate-type-constituents */
193182
onfulfilled?:
194183
| ((value: Response) => TResult1 | PromiseLike<TResult1>)
195184
| undefined
@@ -198,6 +187,7 @@ export class RAGClient {
198187
| ((reason: any) => TResult2 | PromiseLike<TResult2>)
199188
| undefined
200189
| null,
190+
/* eslint-enable @typescript-eslint/no-duplicate-type-constituents */
201191
): Promise<TResult1 | TResult2> {
202192
return fetchRag(
203193
{
@@ -225,7 +215,43 @@ export class RAGClient {
225215
await handleResponseError(response);
226216
}
227217

228-
const data: { data: { embedding: number[] }[] } = await response.json();
229-
return data.data[0].embedding;
218+
const data: unknown = await response.json();
219+
return parseEmbeddingResponse(data);
220+
}
221+
}
222+
223+
function parseResponse(data: unknown): string {
224+
if (
225+
typeof data === "object" &&
226+
data != null &&
227+
"response" in data &&
228+
typeof data.response === "string"
229+
) {
230+
return data.response;
231+
}
232+
233+
throw new Error(
234+
"Expected response to be an object with response key of type string",
235+
);
236+
}
237+
238+
function parseEmbeddingResponse(responseData: unknown): number[] {
239+
if (
240+
typeof responseData === "object" &&
241+
responseData != null &&
242+
"data" in responseData &&
243+
Array.isArray(responseData.data)
244+
) {
245+
const firstItem: unknown = responseData.data[0];
246+
if (
247+
typeof firstItem === "object" &&
248+
firstItem != null &&
249+
"embedding" in firstItem
250+
) {
251+
return firstItem.embedding as number[];
252+
}
230253
}
254+
throw new Error(
255+
"Expected response to be an object with data key of type array of objects with embedding key of type number[]",
256+
);
231257
}

0 commit comments

Comments
 (0)