Add mapping of search results

Author: Dan Jones

src/query/agent.test.ts

@@ -105,9 +105,9 @@ it("search-only mode success: caches searches and sends on subsequent request",
     }),
   } as unknown as WeaviateClient;
 
-  const capturedBodies: ApiSearchModeResponse<undefined>[] = [];
+  const capturedBodies: ApiSearchModeResponse[] = [];
 
-  const apiSuccess: ApiSearchModeResponse<undefined> = {
+  const apiSuccess: ApiSearchModeResponse = {
     original_query: "Test this search only mode!",
     searches: [
       {
@@ -138,23 +138,43 @@ it("search-only mode success: caches searches and sends on subsequent request",
       objects: [
         {
           uuid: "e6dc0a31-76f8-4bd3-b563-677ced6eb557",
-          metadata: {},
-          references: {},
-          vectors: {},
+          metadata: {
+            creation_time: null,
+            update_time: null,
+            distance: null,
+            certainty: null,
+            score: 0.8,
+            explain_score: null,
+            rerank_score: null,
+            is_consistent: null,
+          },
+          references: null,
+          vector: {},
           properties: {
             test_property: 1.0,
             text: "hello",
           },
+          collection: "test_collection",
         },
         {
           uuid: "cf5401cc-f4f1-4eb9-a6a1-173d34f94339",
-          metadata: {},
-          references: {},
-          vectors: {},
+          metadata: {
+            creation_time: null,
+            update_time: null,
+            distance: null,
+            certainty: null,
+            score: 0.5,
+            explain_score: null,
+            rerank_score: null,
+            is_consistent: null,
+          },
+          references: null,
+          vector: {},
           properties: {
             test_property: 2.0,
             text: "world!",
           },
+          collection: "test_collection",
         },
       ],
     },
@@ -164,7 +184,7 @@ it("search-only mode success: caches searches and sends on subsequent request",
   global.fetch = jest.fn((url, init?: RequestInit) => {
     if (init && init.body) {
       capturedBodies.push(
-        JSON.parse(init.body as string) as ApiSearchModeResponse<undefined>,
+        JSON.parse(init.body as string) as ApiSearchModeResponse,
       );
     }
     return Promise.resolve({
@@ -206,7 +226,34 @@ it("search-only mode success: caches searches and sends on subsequent request",
       details: undefined,
     },
     totalTime: 1.5,
-    searchResults: apiSuccess.search_results,
+    searchResults: {
+      objects: [
+        {
+          uuid: "e6dc0a31-76f8-4bd3-b563-677ced6eb557",
+          metadata: {
+            score: 0.8,
+          },
+          vectors: {},
+          properties: {
+            test_property: 1.0,
+            text: "hello",
+          },
+          collection: "test_collection",
+        },
+        {
+          uuid: "cf5401cc-f4f1-4eb9-a6a1-173d34f94339",
+          metadata: {
+            score: 0.5,
+          },
+          vectors: {},
+          properties: {
+            test_property: 2.0,
+            text: "world!",
+          },
+          collection: "test_collection",
+        },
+      ],
+    },
   });
   expect(typeof first.next).toBe("function");
 

src/query/agent.ts

@@ -195,11 +195,11 @@ export class QueryAgent {
    * The returned response includes a `next` method for pagination which
    * reuses the same underlying searches to ensure consistency across pages.
    */
-  async search<T = undefined>(
+  async search(
     query: string,
     { limit = 20, collections }: QueryAgentSearchOnlyOptions = {},
-  ): Promise<SearchModeResponse<T>> {
-    const searcher = new QueryAgentSearcher<T>(this.client, query, {
+  ): Promise<SearchModeResponse> {
+    const searcher = new QueryAgentSearcher(this.client, query, {
       collections: collections ?? this.collections,
       systemPrompt: this.systemPrompt,
       agentsHost: this.agentsHost,

src/query/response/api-response.ts

@@ -1,4 +1,4 @@
-import { WeaviateReturn } from "weaviate-client";
+import { Vectors, WeaviateField } from "weaviate-client";
 
 import {
   NumericMetrics,
@@ -180,10 +180,40 @@ export type ApiSource = {
   collection: string;
 };
 
-export type ApiSearchModeResponse<T> = {
+export type ApiReturnMetadata = {
+  creation_time: Date | null;
+  update_time: Date | null;
+  distance: number | null;
+  certainty: number | null;
+  score: number | null;
+  explain_score: string | null;
+  rerank_score: number | null;
+  is_consistent: boolean | null;
+};
+
+export type ApiWeaviateObject = {
+  /** The returned properties of the object as untyped key-value pairs from the API. */
+  properties: Record<string, WeaviateField>;
+  /** The returned metadata of the object. */
+  metadata: ApiReturnMetadata;
+  /** The returned references of the object. */
+  references: null; // TODO: QA never requests references, so they're never returned?? Check this
+  /** The UUID of the object. */
+  uuid: string;
+  /** The returned vectors of the object. */
+  vector: Vectors; // TODO: note no s!
+  collection: string; // NOTE: NEW
+};
+
+export type ApiWeaviateReturn = {
+  /** The objects that were found by the query. */
+  objects: ApiWeaviateObject[];
+};
+
+export type ApiSearchModeResponse = {
   original_query: string;
   searches?: ApiSearchResult[];
   usage: ApiUsage;
   total_time: number;
-  search_results: WeaviateReturn<T>;
+  search_results: ApiWeaviateReturn;
 };

src/query/response/response-mapping.ts

@@ -1,3 +1,5 @@
+import { ReturnMetadata } from "weaviate-client";
+
 import {
   QueryAgentResponse,
   SearchResult,
@@ -10,6 +12,8 @@ import {
   ProgressMessage,
   DateFilterValue,
   MappedSearchModeResponse,
+  WeaviateObjectWithCollection,
+  WeaviateReturnWithCollection,
 } from "./response.js";
 
 import {
@@ -22,6 +26,8 @@ import {
   ApiSource,
   ApiDateFilterValue,
   ApiSearchModeResponse,
+  ApiWeaviateObject,
+  ApiWeaviateReturn,
 } from "./api-response.js";
 
 import { ServerSentEvent } from "./server-sent-events.js";
@@ -302,19 +308,68 @@ export const mapResponseFromSSE = (
   };
 };
 
-export const mapSearchOnlyResponse = <T>(
-  response: ApiSearchModeResponse<T>,
+const mapWeaviateObject = (
+  object: ApiWeaviateObject,
+): WeaviateObjectWithCollection => {
+  const metadata: ReturnMetadata = {
+    creationTime:
+      object.metadata.creation_time !== null
+        ? object.metadata.creation_time
+        : undefined,
+    updateTime:
+      object.metadata.update_time !== null
+        ? object.metadata.update_time
+        : undefined,
+    distance:
+      object.metadata.distance !== null ? object.metadata.distance : undefined,
+    certainty:
+      object.metadata.certainty !== null
+        ? object.metadata.certainty
+        : undefined,
+    score: object.metadata.score !== null ? object.metadata.score : undefined,
+    explainScore:
+      object.metadata.explain_score !== null
+        ? object.metadata.explain_score
+        : undefined,
+    rerankScore:
+      object.metadata.rerank_score !== null
+        ? object.metadata.rerank_score
+        : undefined,
+    isConsistent:
+      object.metadata.is_consistent !== null
+        ? object.metadata.is_consistent
+        : undefined,
+  };
+
+  return {
+    properties: object.properties,
+    metadata: metadata,
+    references: undefined,
+    uuid: object.uuid,
+    vectors: object.vector,
+    collection: object.collection,
+  };
+};
+
+export const mapWeviateSearchResults = (
+  response: ApiWeaviateReturn,
+): WeaviateReturnWithCollection => ({
+  objects: response.objects.map(mapWeaviateObject),
+});
+
+export const mapSearchOnlyResponse = (
+  response: ApiSearchModeResponse,
 ): {
-  mappedResponse: MappedSearchModeResponse<T>;
+  mappedResponse: MappedSearchModeResponse;
   apiSearches: ApiSearchResult[] | undefined;
 } => {
   const apiSearches = response.searches;
-  const mappedResponse: MappedSearchModeResponse<T> = {
+  const mappedResponse: MappedSearchModeResponse = {
     originalQuery: response.original_query,
     searches: apiSearches ? mapInnerSearches(apiSearches) : undefined,
     usage: mapUsage(response.usage),
     totalTime: response.total_time,
-    searchResults: response.search_results,
+    searchResults: mapWeviateSearchResults(response.search_results),
   };
   return { mappedResponse, apiSearches };
 };

src/query/response/response.ts

@@ -1,4 +1,4 @@
-import { WeaviateReturn } from "weaviate-client";
+import { WeaviateReturn, WeaviateObject } from "weaviate-client";
 
 export type QueryAgentResponse = {
   outputType: "finalState";
@@ -263,12 +263,20 @@ export type StreamedTokens = {
   delta: string;
 };
 
-export type MappedSearchModeResponse<T> = {
+export type WeaviateObjectWithCollection = WeaviateObject<undefined> & {
+  collection: string;
+};
+
+export type WeaviateReturnWithCollection = WeaviateReturn<undefined> & {
+  objects: WeaviateObjectWithCollection[];
+};
+
+export type MappedSearchModeResponse = {
   originalQuery: string;
   searches?: SearchResult[];
   usage: Usage;
   totalTime: number;
-  searchResults: WeaviateReturn<T>;
+  searchResults: WeaviateReturnWithCollection;
 };
 
 /** Options for the executing a prepared QueryAgent search. */
@@ -279,6 +287,6 @@ export type SearchExecutionOptions = {
   offset?: number;
 };
 
-export type SearchModeResponse<T> = MappedSearchModeResponse<T> & {
-  next: (options: SearchExecutionOptions) => Promise<SearchModeResponse<T>>;
+export type SearchModeResponse = MappedSearchModeResponse & {
+  next: (options: SearchExecutionOptions) => Promise<SearchModeResponse>;
 };

src/query/search.ts

@@ -26,7 +26,7 @@ import {
  *
  * For more information, see the [Weaviate Query Agent Docs](https://weaviate.io/developers/agents/query)
  */
-export class QueryAgentSearcher<T> {
+export class QueryAgentSearcher {
   private agentsHost: string;
   private query: string;
   private collections: (string | QueryAgentCollectionConfig)[];
@@ -103,9 +103,10 @@ export class QueryAgentSearcher<T> {
    * @param [options.offset] - The offset to start from. If not specified, retrieval begins from the first object.
    * @returns A SearchModeResponse object containing the results, usage, and underlying searches performed.
    */
-  async run({ limit = 20, offset = 0 }: SearchExecutionOptions = {}): Promise<
-    SearchModeResponse<T>
-  > {
+  async run({
+    limit = 20,
+    offset = 0,
+  }: SearchExecutionOptions = {}): Promise<SearchModeResponse> {
     if (!this.collections || this.collections.length === 0) {
       throw Error("No collections provided to the query agent.");
     }
@@ -121,9 +122,9 @@ export class QueryAgentSearcher<T> {
     if (!response.ok) {
       await handleError(await response.text());
     }
-    const parsedResponse = (await response.json()) as ApiSearchModeResponse<T>;
+    const parsedResponse = (await response.json()) as ApiSearchModeResponse;
     const { mappedResponse, apiSearches } =
-      mapSearchOnlyResponse<T>(parsedResponse);
+      mapSearchOnlyResponse(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