Make searcher generic to support generic WeaviateReturn

Author: Dan Jones

src/query/agent.ts

@@ -194,10 +194,10 @@ export class QueryAgent {
    * @param options - Additional options for the searcher.
    * @returns The searcher for the query agent.
    */
-  prepareSearch(
+  prepareSearch<T = undefined>(
     query: string,
     { collections }: QueryAgentSearchOnlyOptions = {},
-  ): QueryAgentSearcher {
+  ): QueryAgentSearcher<T> {
     return new QueryAgentSearcher(this.client, query, {
       collections: collections ?? this.collections,
       systemPrompt: this.systemPrompt,

src/query/response/api-response.ts

@@ -180,10 +180,10 @@ export type ApiSource = {
   collection: string;
 };
 
-export type ApiSearchModeResponse = {
+export type ApiSearchModeResponse<T> = {
   original_query: string;
   searches?: ApiSearchResult[];
   usage: ApiUsage;
   total_time: number;
-  search_results: WeaviateReturn<undefined>;
+  search_results: WeaviateReturn<T>;
 };

src/query/response/response-mapping.ts

@@ -302,14 +302,14 @@ export const mapResponseFromSSE = (
   };
 };
 
-export const mapSearchOnlyResponse = (
-  response: ApiSearchModeResponse,
+export const mapSearchOnlyResponse = <T>(
+  response: ApiSearchModeResponse<T>,
 ): {
-  mappedResponse: SearchModeResponse;
+  mappedResponse: SearchModeResponse<T>;
   apiSearches: ApiSearchResult[] | undefined;
 } => {
   const apiSearches = response.searches;
-  const mappedResponse: SearchModeResponse = {
+  const mappedResponse: SearchModeResponse<T> = {
     originalQuery: response.original_query,
     searches: apiSearches ? mapInnerSearches(apiSearches) : undefined,
     usage: mapUsage(response.usage),

src/query/response/response.ts

@@ -263,10 +263,10 @@ export type StreamedTokens = {
   delta: string;
 };
 
-export type SearchModeResponse = {
+export type SearchModeResponse<T> = {
   originalQuery: string;
   searches?: SearchResult[];
   usage: Usage;
   totalTime: number;
-  searchResults: WeaviateReturn<undefined>;
+  searchResults: WeaviateReturn<T>;
 };

src/query/search.ts

@@ -17,7 +17,7 @@ import {
  *
  * For more information, see the [Weaviate Query Agent Docs](https://weaviate.io/developers/agents/query)
  */
-export class QueryAgentSearcher {
+export class QueryAgentSearcher<T> {
   //private headers: Record<string, string> = {};
   //private connectionHeaders: HeadersInit | undefined;
   private agentsHost: string;
@@ -84,7 +84,7 @@ export class QueryAgentSearcher {
     };
   }
 
-  async execute(options: SearchExecutionOptions): Promise<SearchModeResponse> {
+  async execute(options: SearchExecutionOptions): Promise<SearchModeResponse<T>> {
     if (!this.collections || this.collections.length === 0) {
       throw Error("No collections provided to the query agent.");
     }
@@ -101,9 +101,9 @@ export class QueryAgentSearcher {
     if (!response.ok) {
       await handleError(await response.text());
     }
-    const parsedResponse = (await response.json()) as ApiSearchModeResponse;
+    const parsedResponse = (await response.json()) as ApiSearchModeResponse<T>;
     const { mappedResponse, apiSearches } =
-      mapSearchOnlyResponse(parsedResponse);
+      mapSearchOnlyResponse<T>(parsedResponse);
     // If we successfully mapped the searches, cache them for the next request.
     // Since this cache is a private internal value, there's not point in mapping
     // back and forth between the exported and API types, so we cache apiSearches