weaviate-agents
Version:
JS/TS client for Weaviate Agents
80 lines (79 loc) • 4.18 kB
JavaScript
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
import { mapSearchOnlyResponse } from "./response/response-mapping.js";
import { mapCollections } from "./collection.js";
import { handleError } from "./response/error.js";
import { getHeaders } from "./connection.js";
/**
* A configured searcher for the QueryAgent.
*
* This is used internally by the QueryAgent class to run search-mode queries.
* After the first request is made, the underlying searches are cached and can
* be reused for paginating through the a consistent set of results.
*/
export class QueryAgentSearcher {
constructor(client, query, collections, systemPrompt, agentsHost) {
this.client = client;
this.query = query;
this.collections = collections;
this.systemPrompt = systemPrompt;
this.agentsHost = agentsHost;
}
buildRequestBody(limit, offset, connectionHeaders) {
const base = {
headers: connectionHeaders,
original_query: typeof this.query === "string" ? this.query : { messages: this.query },
collections: mapCollections(this.collections),
limit,
offset,
};
if (this.cachedSearches === undefined) {
return Object.assign(Object.assign({}, base), { searches: null, system_prompt: this.systemPrompt || null });
}
return Object.assign(Object.assign({}, base), { searches: this.cachedSearches });
}
/**
* Run the search-only agent with the given limit and offset values.
*
* Calling this method multiple times on the same QueryAgentSearcher instance will result
* in the same underlying searches being performed each time, allowing you to paginate
* over a consistent results set.
*
* @param [options] - Options for executing the search
* @param [options.limit] - The maximum number of results to return. Defaults to 20 if not specified.
* @param [options.offset] - The offset to start from.
* @returns A SearchModeResponse object containing the results, usage, and underlying searches performed.
*/
run(_a) {
return __awaiter(this, arguments, void 0, function* ({ limit = 20, offset, }) {
if (!this.collections || this.collections.length === 0) {
throw Error("No collections provided to the query agent.");
}
const { requestHeaders, connectionHeaders } = yield getHeaders(this.client);
const response = yield fetch(`${this.agentsHost}/query/search_only`, {
method: "POST",
headers: requestHeaders,
body: JSON.stringify(this.buildRequestBody(limit, offset, connectionHeaders)),
});
if (!response.ok) {
yield handleError(yield response.text());
}
const parsedResponse = (yield response.json());
const { mappedResponse, apiSearches } = 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
if (mappedResponse.searches) {
this.cachedSearches = apiSearches;
}
return Object.assign(Object.assign({}, mappedResponse), { next: (options) => __awaiter(this, void 0, void 0, function* () { return this.run(options); }) });
});
}
}