Skip to content

Commit c183f83

Browse files
committed
introduce ask/askStream and deprecate run/stream
1 parent 820e581 commit c183f83

File tree

1 file changed

+143
-0
lines changed

1 file changed

+143
-0
lines changed

src/query/agent.ts

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ export class QueryAgent {
5353
/**
5454
* Run the query agent.
5555
*
56+
* @deprecated Use {@link ask} instead.
5657
* @param query - The natural language query string for the agent.
5758
* @param options - Additional options for the run.
5859
* @returns The response from the query agent.
@@ -93,9 +94,53 @@ export class QueryAgent {
9394
return mapResponse(await response.json());
9495
}
9596

97+
/**
98+
* Ask query agent a question.
99+
*
100+
* @param query - The natural language query string for the agent.
101+
* @param options - Additional options for the run.
102+
* @returns The response from the query agent.
103+
*/
104+
async ask(
105+
query: string,
106+
{ collections, context }: QueryAgentRunOptions = {},
107+
): Promise<QueryAgentResponse> {
108+
const targetCollections = collections ?? this.collections;
109+
if (!targetCollections) {
110+
throw Error("No collections provided to the query agent.");
111+
}
112+
113+
const { host, bearerToken, headers } =
114+
await this.client.getConnectionDetails();
115+
116+
const response = await fetch(`${this.agentsHost}/agent/query`, {
117+
method: "POST",
118+
headers: {
119+
"Content-Type": "application/json",
120+
Authorization: bearerToken!,
121+
"X-Weaviate-Cluster-Url": host,
122+
"X-Agent-Request-Origin": "typescript-client",
123+
},
124+
body: JSON.stringify({
125+
headers,
126+
query,
127+
collections: mapCollections(targetCollections),
128+
system_prompt: this.systemPrompt,
129+
previous_response: context ? mapApiResponse(context) : undefined,
130+
}),
131+
});
132+
133+
if (!response.ok) {
134+
await handleError(await response.text());
135+
}
136+
137+
return mapResponse(await response.json());
138+
}
139+
96140
/**
97141
* Stream responses from the query agent.
98142
*
143+
* @deprecated Use {@link askStream} instead.
99144
* @param query - The natural language query string for the agent.
100145
* @param options - Additional options for the run.
101146
* @returns The response from the query agent.
@@ -107,20 +152,23 @@ export class QueryAgent {
107152
includeFinalState: false;
108153
},
109154
): AsyncGenerator<StreamedTokens>;
155+
/** @deprecated Use {@link askStream} instead. */
110156
stream(
111157
query: string,
112158
options: QueryAgentStreamOptions & {
113159
includeProgress: false;
114160
includeFinalState?: true;
115161
},
116162
): AsyncGenerator<StreamedTokens | QueryAgentResponse>;
163+
/** @deprecated Use {@link askStream} instead. */
117164
stream(
118165
query: string,
119166
options: QueryAgentStreamOptions & {
120167
includeProgress?: true;
121168
includeFinalState: false;
122169
},
123170
): AsyncGenerator<ProgressMessage | StreamedTokens>;
171+
/** @deprecated Use {@link askStream} instead. */
124172
stream(
125173
query: string,
126174
options?: QueryAgentStreamOptions & {
@@ -188,6 +236,101 @@ export class QueryAgent {
188236
}
189237
}
190238

239+
/**
240+
* Ask query agent a question and stream the response.
241+
*
242+
* @param query - The natural language query string for the agent.
243+
* @param options - Additional options for the run.
244+
* @returns The response from the query agent.
245+
*/
246+
askStream(
247+
query: string,
248+
options: QueryAgentStreamOptions & {
249+
includeProgress: false;
250+
includeFinalState: false;
251+
},
252+
): AsyncGenerator<StreamedTokens>;
253+
askStream(
254+
query: string,
255+
options: QueryAgentStreamOptions & {
256+
includeProgress: false;
257+
includeFinalState?: true;
258+
},
259+
): AsyncGenerator<StreamedTokens | QueryAgentResponse>;
260+
askStream(
261+
query: string,
262+
options: QueryAgentStreamOptions & {
263+
includeProgress?: true;
264+
includeFinalState: false;
265+
},
266+
): AsyncGenerator<ProgressMessage | StreamedTokens>;
267+
askStream(
268+
query: string,
269+
options?: QueryAgentStreamOptions & {
270+
includeProgress?: true;
271+
includeFinalState?: true;
272+
},
273+
): AsyncGenerator<ProgressMessage | StreamedTokens | QueryAgentResponse>;
274+
async *askStream(
275+
query: string,
276+
{
277+
collections,
278+
context,
279+
includeProgress,
280+
includeFinalState,
281+
}: QueryAgentStreamOptions = {},
282+
): AsyncGenerator<ProgressMessage | StreamedTokens | QueryAgentResponse> {
283+
const targetCollections = collections ?? this.collections;
284+
285+
if (!targetCollections) {
286+
throw Error("No collections provided to the query agent.");
287+
}
288+
289+
const { host, bearerToken, headers } =
290+
await this.client.getConnectionDetails();
291+
292+
const sseStream = fetchServerSentEvents(
293+
`${this.agentsHost}/agent/stream_query`,
294+
{
295+
method: "POST",
296+
headers: {
297+
"Content-Type": "application/json",
298+
Authorization: bearerToken!,
299+
"X-Weaviate-Cluster-Url": host,
300+
"X-Agent-Request-Origin": "typescript-client",
301+
},
302+
body: JSON.stringify({
303+
headers,
304+
query,
305+
collections: mapCollections(targetCollections),
306+
system_prompt: this.systemPrompt,
307+
previous_response: context ? mapApiResponse(context) : undefined,
308+
include_progress: includeProgress ?? true,
309+
include_final_state: includeFinalState ?? true,
310+
}),
311+
},
312+
);
313+
314+
for await (const event of sseStream) {
315+
if (event.event === "error") {
316+
await handleError(event.data);
317+
}
318+
319+
let output: ProgressMessage | StreamedTokens | QueryAgentResponse;
320+
if (event.event === "progress_message") {
321+
output = mapProgressMessageFromSSE(event);
322+
} else if (event.event === "streamed_tokens") {
323+
output = mapStreamedTokensFromSSE(event);
324+
} else if (event.event === "final_state") {
325+
output = mapResponseFromSSE(event);
326+
} else {
327+
throw new Error(`Unexpected event type: ${event.event}: ${event.data}`);
328+
}
329+
330+
yield output;
331+
}
332+
}
333+
191334
/**
192335
* Run the Query Agent search-only mode.
193336
*

0 commit comments

Comments
 (0)