swith search mode to v2 API

Author: augustas1

src/query/agent.test.ts

@@ -108,30 +108,22 @@ it("search-only mode success: caches searches and sends on subsequent request",
   const capturedBodies: ApiSearchModeResponse[] = [];
 
   const apiSuccess: ApiSearchModeResponse = {
-    original_query: "Test this search only mode!",
     searches: [
       {
-        queries: ["search query"],
-        filters: [
-          [
-            {
-              filter_type: "integer",
-              property_name: "test_property",
-              operator: ComparisonOperator.GreaterThan,
-              value: 0,
-            },
-          ],
-        ],
-        filter_operators: "AND",
+        query: "search query",
+        filters: {
+          filter_type: "integer",
+          property_name: "test_property",
+          operator: ComparisonOperator.GreaterThan,
+          value: 0,
+        },
         collection: "test_collection",
       },
     ],
     usage: {
-      requests: 0,
-      request_tokens: undefined,
-      response_tokens: undefined,
-      total_tokens: undefined,
-      details: undefined,
+      model_units: 1,
+      usage_in_plan: true,
+      remaining_plan_requests: 2,
     },
     total_time: 1.5,
     search_results: {
@@ -200,30 +192,22 @@ it("search-only mode success: caches searches and sends on subsequent request",
     collections: ["test_collection"],
   });
   expect(first).toMatchObject({
-    originalQuery: apiSuccess.original_query,
     searches: [
       {
         collection: "test_collection",
-        queries: ["search query"],
-        filters: [
-          [
-            {
-              filterType: "integer",
-              propertyName: "test_property",
-              operator: ComparisonOperator.GreaterThan,
-              value: 0,
-            },
-          ],
-        ],
-        filterOperators: "AND",
+        query: "search query",
+        filters: {
+          filterType: "integer",
+          propertyName: "test_property",
+          operator: ComparisonOperator.GreaterThan,
+          value: 0,
+        },
       },
     ],
     usage: {
-      requests: 0,
-      requestTokens: undefined,
-      responseTokens: undefined,
-      totalTokens: undefined,
-      details: undefined,
+      modelUnits: 1,
+      usageInPlan: true,
+      remainingPlanRequests: 2,
     },
     totalTime: 1.5,
     searchResults: {
@@ -266,7 +250,6 @@ it("search-only mode success: caches searches and sends on subsequent request",
   expect(capturedBodies[1].searches).toEqual(apiSuccess.searches);
   // Response mapping should be the same (because response is mocked)
   expect(second).toMatchObject({
-    originalQuery: apiSuccess.original_query,
     searches: first.searches,
     usage: first.usage,
     totalTime: first.totalTime,

src/query/agent.ts

@@ -9,7 +9,6 @@ import {
   mapResponse,
   mapProgressMessageFromSSE,
   mapStreamedTokensFromSSE,
-  mapResponseFromSSE,
   mapAskModeResponse,
 } from "./response/response-mapping.js";
 import { mapApiResponse } from "./response/api-response-mapping.js";

src/query/response/api-response.ts

@@ -247,9 +247,8 @@ export type ApiWeaviateReturn = {
 };
 
 export type ApiSearchModeResponse = {
-  original_query: string;
-  searches?: ApiSearchResult[];
-  usage: ApiUsage;
+  searches?: ApiSearch[];
+  usage: ApiModelUnitUsage;
   total_time: number;
   search_results: ApiWeaviateReturn;
 };

src/query/response/response-mapping.ts

@@ -2,11 +2,9 @@ import { ReturnMetadata } from "weaviate-client";
 
 import {
   QueryAgentResponse,
-  SearchResult,
   PropertyFilter,
   AggregationResult,
   PropertyAggregation,
-  Usage,
   Source,
   StreamedTokens,
   ProgressMessage,
@@ -16,22 +14,24 @@ import {
   SearchModeResponse,
   FilterAndOr,
   AskModeResponse,
+  Search,
+  ModelUnitUsage,
 } from "./response.js";
 
 import {
   ApiQueryAgentResponse,
-  ApiSearchResult,
   ApiPropertyFilter,
   ApiAggregationResult,
   ApiPropertyAggregation,
-  ApiUsage,
   ApiSource,
   ApiDateFilterValue,
   ApiSearchModeResponse,
   ApiWeaviateObject,
   ApiWeaviateReturn,
   ApiFilterAndOr,
   ApiAskModeResponse,
+  ApiSearch,
+  ApiModelUnitUsage,
 } from "./api-response.js";
 
 import { ServerSentEvent } from "./server-sent-events.js";
@@ -41,22 +41,14 @@ export const mapAskModeResponse = (
 ): AskModeResponse => {
   const properties: AskModeResponseProperties = {
     outputType: "finalState",
-    searches: response.searches.map((search) => ({
-      query: search.query,
-      filters: search.filters ? mapFilter(search.filters) : undefined,
-      collection: search.collection,
-    })),
+    searches: mapSearches(response.searches),
     aggregations: response.aggregations.map((aggregation) => ({
       groupbyProperty: aggregation.groupby_property,
       aggregation: mapPropertyAggregation(aggregation.aggregation),
       filters: aggregation.filters ? mapFilter(aggregation.filters) : undefined,
       collection: aggregation.collection,
     })),
-    usage: {
-      modelUnits: response.usage.model_units,
-      usageInPlan: response.usage.usage_in_plan,
-      remainingPlanRequests: response.usage.remaining_plan_requests,
-    },
+    usage: mapUsage(response.usage),
     totalTime: response.total_time,
     isPartialAnswer: response.is_partial_answer,
     missingInformation: response.missing_information,
@@ -70,16 +62,42 @@ export const mapAskModeResponse = (
   };
 };
 
+const mapSearches = (searches: ApiSearch[]): Search[] =>
+  searches.map((search) => ({
+    query: search.query,
+    filters: search.filters ? mapFilter(search.filters) : undefined,
+    collection: search.collection,
+  }));
+
+const mapUsage = (usage: ApiModelUnitUsage): ModelUnitUsage => ({
+  modelUnits: usage.model_units,
+  usageInPlan: usage.usage_in_plan,
+  remainingPlanRequests: usage.remaining_plan_requests,
+});
+
 export const mapResponse = (
   response: ApiQueryAgentResponse,
 ): QueryAgentResponse => {
   const properties: ResponseProperties = {
     outputType: "finalState",
     originalQuery: response.original_query,
     collectionNames: response.collection_names,
-    searches: mapSearches(response.searches),
+    searches: response.searches.map((searches) =>
+      searches.map((result) => ({
+        collection: result.collection,
+        queries: result.queries,
+        filters: result.filters.map((filter) => filter.map(mapPropertyFilter)),
+        filterOperators: result.filter_operators,
+      })),
+    ),
     aggregations: mapAggregations(response.aggregations),
-    usage: mapUsage(response.usage),
+    usage: {
+      requests: response.usage.requests,
+      requestTokens: response.usage.request_tokens,
+      responseTokens: response.usage.response_tokens,
+      totalTokens: response.usage.total_tokens,
+      details: response.usage.details,
+    },
     totalTime: response.total_time,
     isPartialAnswer: response.is_partial_answer,
     missingInformation: response.missing_information,
@@ -93,17 +111,6 @@ export const mapResponse = (
   };
 };
 
-const mapInnerSearches = (searches: ApiSearchResult[]): SearchResult[] =>
-  searches.map((result) => ({
-    collection: result.collection,
-    queries: result.queries,
-    filters: result.filters.map((filter) => filter.map(mapPropertyFilter)),
-    filterOperators: result.filter_operators,
-  }));
-
-const mapSearches = (searches: ApiSearchResult[][]): SearchResult[][] =>
-  searches.map((searchGroup) => mapInnerSearches(searchGroup));
-
 const mapDatePropertyFilter = (
   filterValue: ApiDateFilterValue,
 ): DateFilterValue | undefined => {
@@ -269,14 +276,6 @@ const mapPropertyAggregation = (
       : undefined,
 });
 
-const mapUsage = (usage: ApiUsage): Usage => ({
-  requests: usage.requests,
-  requestTokens: usage.request_tokens,
-  responseTokens: usage.response_tokens,
-  totalTokens: usage.total_tokens,
-  details: usage.details,
-});
-
 const mapSources = (sources: ApiSource[]): Source[] =>
   sources.map((source) => ({
     objectId: source.object_id,
@@ -332,31 +331,6 @@ export const mapStreamedTokensFromSSE = (
   };
 };
 
-export const mapResponseFromSSE = (
-  sse: ServerSentEvent,
-): QueryAgentResponse => {
-  const data: ApiQueryAgentResponse = JSON.parse(sse.data);
-
-  const properties: ResponseProperties = {
-    outputType: "finalState",
-    originalQuery: data.original_query,
-    collectionNames: data.collection_names,
-    searches: mapSearches(data.searches),
-    aggregations: mapAggregations(data.aggregations),
-    usage: mapUsage(data.usage),
-    totalTime: data.total_time,
-    isPartialAnswer: data.is_partial_answer,
-    missingInformation: data.missing_information,
-    finalAnswer: data.final_answer,
-    sources: mapSources(data.sources),
-  };
-
-  return {
-    ...properties,
-    display: () => display(properties),
-  };
-};
-
 const mapWeaviateObject = (
   object: ApiWeaviateObject,
 ): WeaviateObjectWithCollection => {
@@ -410,12 +384,11 @@ export const mapSearchOnlyResponse = (
   response: ApiSearchModeResponse,
 ): {
   mappedResponse: Omit<SearchModeResponse, "next">;
-  apiSearches: ApiSearchResult[] | undefined;
+  apiSearches: ApiSearch[] | undefined;
 } => {
   const apiSearches = response.searches;
   const mappedResponse: Omit<SearchModeResponse, "next"> = {
-    originalQuery: response.original_query,
-    searches: apiSearches ? mapInnerSearches(apiSearches) : undefined,
+    searches: apiSearches ? mapSearches(apiSearches) : undefined,
     usage: mapUsage(response.usage),
     totalTime: response.total_time,
     searchResults: mapWeviateSearchResults(response.search_results),

src/query/response/response.ts

@@ -317,9 +317,8 @@ export type SearchExecutionOptions = {
 };
 
 export type SearchModeResponse = {
-  originalQuery: string;
-  searches?: SearchResult[];
-  usage: Usage;
+  searches?: Search[];
+  usage: ModelUnitUsage;
   totalTime: number;
   searchResults: WeaviateReturnWithCollection;
   next: (options: SearchExecutionOptions) => Promise<SearchModeResponse>;

src/query/search.ts

@@ -6,10 +6,7 @@ import {
 import { mapSearchOnlyResponse } from "./response/response-mapping.js";
 import { mapCollections } from "./collection.js";
 import { handleError } from "./response/error.js";
-import {
-  ApiSearchModeResponse,
-  ApiSearchResult,
-} from "./response/api-response.js";
+import { ApiSearch, ApiSearchModeResponse } from "./response/api-response.js";
 import { QueryAgentQuery } from "./agent.js";
 import { QueryAgentCollection } from "./collection.js";
 import { getHeaders } from "./connection.js";
@@ -28,7 +25,7 @@ import { getHeaders } from "./connection.js";
  * For more information, see the [Weaviate Query Agent Docs](https://weaviate.io/developers/agents/query)
  */
 export class QueryAgentSearcher {
-  private cachedSearches?: ApiSearchResult[];
+  private cachedSearches?: ApiSearch[];
 
   constructor(
     private client: WeaviateClient,
@@ -85,7 +82,7 @@ export class QueryAgentSearcher {
     }
     const { requestHeaders, connectionHeaders } = await getHeaders(this.client);
 
-    const response = await fetch(`${this.agentsHost}/agent/search_only`, {
+    const response = await fetch(`${this.agentsHost}/query/search_only`, {
       method: "POST",
       headers: requestHeaders,
       body: JSON.stringify(