Update tests

Author: Dan Jones

src/query/agent.test.ts

@@ -1,7 +1,14 @@
 import { WeaviateClient } from "weaviate-client";
 import { QueryAgent } from "./agent.js";
 import { ApiQueryAgentResponse } from "./response/api-response.js";
-import { QueryAgentResponse } from "./response/response.js";
+import {
+  QueryAgentResponse,
+  ComparisonOperator,
+  SearchModeResponse,
+} from "./response/response.js";
+import { QueryAgentSearcher } from "./search.js";
+import { ApiSearchModeResponse } from "./response/api-response.js";
+import { QueryAgentError } from "./response/error.js";
 
 it("runs the query agent", async () => {
   const mockClient = {
@@ -93,3 +100,184 @@ it("runs the query agent", async () => {
     display: expect.any(Function),
   });
 });
+
+it("prepareSearch returns a QueryAgentSearcher", async () => {
+  const mockClient = {
+    getConnectionDetails: jest.fn().mockResolvedValue({
+      host: "test-cluster",
+      bearerToken: "test-token",
+      headers: { "X-Provider": "test-key" },
+    }),
+  } as unknown as WeaviateClient;
+
+  const agent = new QueryAgent(mockClient, {
+    systemPrompt: "test system prompt",
+  });
+
+  const searcher = agent.prepareSearch("test query");
+  expect(searcher).toBeInstanceOf(QueryAgentSearcher);
+});
+
+it("search-only mode success: caches searches and sends on subsequent request", async () => {
+  const mockClient = {
+    getConnectionDetails: jest.fn().mockResolvedValue({
+      host: "test-cluster",
+      bearerToken: "test-token",
+      headers: { "X-Provider": "test-key" },
+    }),
+  } as unknown as WeaviateClient;
+
+  const capturedBodies: ApiSearchModeResponse<undefined>[] = [];
+
+  const apiSuccess: ApiSearchModeResponse<undefined> = {
+    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",
+        collection: "test_collection",
+      },
+    ],
+    usage: {
+      requests: 0,
+      request_tokens: undefined,
+      response_tokens: undefined,
+      total_tokens: undefined,
+      details: undefined,
+    },
+    total_time: 1.5,
+    search_results: {
+      objects: [
+        {
+          uuid: "e6dc0a31-76f8-4bd3-b563-677ced6eb557",
+          metadata: {},
+          references: {},
+          vectors: {},
+          properties: {
+            test_property: 1.0,
+            text: "hello",
+          },
+        },
+        {
+          uuid: "cf5401cc-f4f1-4eb9-a6a1-173d34f94339",
+          metadata: {},
+          references: {},
+          vectors: {},
+          properties: {
+            test_property: 2.0,
+            text: "world!",
+          },
+        },
+      ],
+    },
+  };
+
+  // Mock the API response, and capture the request body to assert later
+  global.fetch = jest.fn((url, init?: RequestInit) => {
+    if (init && init.body) {
+      capturedBodies.push(
+        JSON.parse(init.body as string) as ApiSearchModeResponse<undefined>,
+      );
+    }
+    return Promise.resolve({
+      ok: true,
+      json: () => Promise.resolve(apiSuccess),
+    } as Response);
+  }) as jest.Mock;
+
+  const agent = new QueryAgent(mockClient);
+  const searcher = agent.prepareSearch("test query", {
+    collections: ["test_collection"],
+  });
+
+  const first = await searcher.execute({ limit: 2, offset: 0 });
+
+  expect(first).toEqual<SearchModeResponse<undefined>>({
+    originalQuery: apiSuccess.original_query,
+    searches: [
+      {
+        collection: "test_collection",
+        queries: ["search query"],
+        filters: [
+          [
+            {
+              filterType: "integer",
+              propertyName: "test_property",
+              operator: ComparisonOperator.GreaterThan,
+              value: 0,
+            },
+          ],
+        ],
+        filterOperators: "AND",
+      },
+    ],
+    usage: {
+      requests: 0,
+      requestTokens: undefined,
+      responseTokens: undefined,
+      totalTokens: undefined,
+      details: undefined,
+    },
+    totalTime: 1.5,
+    searchResults: apiSuccess.search_results,
+  });
+
+  // First request should have searches: null (generation request)
+  expect(capturedBodies[0].searches).toBeNull();
+  const second = await searcher.execute({ limit: 2, offset: 1 });
+  // Second request should include the original searches (execution request)
+  expect(capturedBodies[1].searches).toEqual(apiSuccess.searches);
+  // Response mapping should be the same (because response is mocked)
+  expect(second).toEqual(first);
+});
+
+it("search-only mode failure propagates QueryAgentError", async () => {
+  const mockClient = {
+    getConnectionDetails: jest.fn().mockResolvedValue({
+      host: "test-cluster",
+      bearerToken: "test-token",
+      headers: { "X-Provider": "test-key" },
+    }),
+  } as unknown as WeaviateClient;
+
+  const errorJson = {
+    error: {
+      message: "Test error message",
+      code: "test_error_code",
+      details: { info: "test detail" },
+    },
+  };
+
+  global.fetch = jest.fn(() =>
+    Promise.resolve({
+      ok: false,
+      text: () => Promise.resolve(JSON.stringify(errorJson)),
+    } as Response),
+  ) as jest.Mock;
+
+  const agent = new QueryAgent(mockClient);
+  const searcher = agent.prepareSearch("test query", {
+    collections: ["test_collection"],
+  });
+
+  try {
+    await searcher.execute({ limit: 2, offset: 0 });
+  } catch (err) {
+    expect(err).toBeInstanceOf(QueryAgentError);
+    expect(err).toMatchObject({
+      message: "Test error message",
+      code: "test_error_code",
+      details: { info: "test detail" },
+    });
+  }
+});