Context (#5)

* add context

* fix typo

Author: augustas1

src/query/agent.ts

@@ -1,6 +1,7 @@
 import { WeaviateClient } from "weaviate-client";
 import { QueryAgentResponse } from "./response/response.js";
 import { mapResponse } from "./response/response-mapping.js";
+import { mapApiResponse } from "./response/api-reponse-mapping.js";
 
 /**
  * An agent for executing agentic queries against Weaviate.
@@ -43,7 +44,7 @@ export class QueryAgent {
    */
   async run(
     query: string,
-    { viewProperties }: QueryAgentRunOptions = {}
+    { viewProperties, context }: QueryAgentRunOptions = {}
   ): Promise<QueryAgentResponse> {
     const { host, bearerToken, headers } =
       await this.client.getConnectionDetails();
@@ -61,16 +62,15 @@ export class QueryAgent {
         collection_names: this.collections,
         collection_view_properties: viewProperties,
         system_prompt: this.systemPrompt,
+        previous_response: context ? mapApiResponse(context) : undefined,
       }),
     });
 
-    const responseBody = await response.json();
-
     if (!response.ok) {
-      throw Error(`Query agent failed. ${JSON.stringify(responseBody)}`);
+      throw Error(`Query agent failed. ${await response.text()}`);
     }
 
-    return mapResponse(responseBody);
+    return mapResponse(await response.json());
   }
 }
 
@@ -86,4 +86,6 @@ export type QueryAgentOptions = {
 export type QueryAgentRunOptions = {
   /** List of of property names the agent has the ability to view across all collections. */
   viewProperties?: string[];
+  /** Previous response from the agent. */
+  context?: QueryAgentResponse;
 };

src/query/response/api-response-mapping.ts

@@ -0,0 +1,94 @@
+import {
+  AggregationResult,
+  PropertyAggregation,
+  PropertyFilter,
+  QueryAgentResponse,
+  SearchResult,
+  Source,
+  Usage,
+} from "./response.js";
+
+import {
+  ApiAggregationResult,
+  ApiPropertyAggregation,
+  ApiPropertyFilter,
+  ApiQueryAgentResponse,
+  ApiSearchResult,
+  ApiSource,
+  ApiUsage,
+} from "./api-response.js";
+
+export const mapApiResponse = (
+  response: QueryAgentResponse
+): ApiQueryAgentResponse => ({
+  original_query: response.originalQuery,
+  collection_names: response.collectionNames,
+  searches: mapApiSearches(response.searches),
+  aggregations: mapApiAggregations(response.aggregations),
+  usage: mapApiUsage(response.usage),
+  total_time: response.totalTime,
+  aggregation_answer: response.aggregationAnswer,
+  has_aggregation_answer: response.hasAggregationAnswer,
+  has_search_answer: response.hasSearchAnswer,
+  is_partial_answer: response.isPartialAnswer,
+  missing_information: response.missingInformation,
+  final_answer: response.finalAnswer,
+  sources: mapApiSources(response.sources),
+});
+
+const mapApiSearches = (searches: SearchResult[][]): ApiSearchResult[][] =>
+  searches.map((searchGroup) =>
+    searchGroup.map((result) => ({
+      collection: result.collection,
+      queries: result.queries,
+      filters: result.filters.map(mapApiPropertyFilters),
+      filter_operators: result.filterOperators,
+    }))
+  );
+
+const mapApiPropertyFilters = (
+  filters: PropertyFilter[]
+): ApiPropertyFilter[] =>
+  filters.map((filter) => ({
+    property_name: filter.propertyName,
+    operator: filter.operator,
+    value: filter.value,
+  }));
+
+const mapApiAggregations = (
+  aggregations: AggregationResult[][]
+): ApiAggregationResult[][] =>
+  aggregations.map((aggregationGroup) =>
+    aggregationGroup.map((result) => ({
+      collection: result.collection,
+      search_query: result.searchQuery,
+      groupby_property: result.groupbyProperty,
+      aggregations: result.aggregations.map(mapApiPropertyAggregation),
+      filters: mapApiPropertyFilters(result.filters),
+    }))
+  );
+
+const mapApiPropertyAggregation = (
+  aggregation: PropertyAggregation
+): ApiPropertyAggregation => ({
+  property_name: aggregation.propertyName,
+  metrics: aggregation.metrics,
+  top_occurrences_limit:
+    "topOccurrencesLimit" in aggregation
+      ? aggregation.topOccurrencesLimit
+      : undefined,
+});
+
+const mapApiUsage = (usage: Usage): ApiUsage => ({
+  requests: usage.requests,
+  request_tokens: usage.requestTokens,
+  response_tokens: usage.responseTokens,
+  total_tokens: usage.totalTokens,
+  details: usage.details,
+});
+
+const mapApiSources = (sources: Source[]): ApiSource[] =>
+  sources.map((source) => ({
+    object_id: source.objectId,
+    collection: source.collection,
+  }));

src/query/response/api-response.ts

@@ -24,7 +24,7 @@ export type ApiQueryAgentResponse = {
 export type ApiSearchResult = {
   collection: string;
   queries: string[];
-  filters: ApiPropertyFilter[];
+  filters: ApiPropertyFilter[][];
   filter_operators: "AND" | "OR";
 };
 

src/query/response/response-mapping.ts

@@ -48,27 +48,28 @@ const mapSearches = (searches: ApiSearchResult[][]): SearchResult[][] =>
     searchGroup.map((result) => ({
       collection: result.collection,
       queries: result.queries,
-      filters: result.filters.map(mapPropertyFilter),
+      filters: result.filters.map(mapPropertyFilters),
       filterOperators: result.filter_operators,
     }))
   );
 
-const mapPropertyFilter = (filter: ApiPropertyFilter): PropertyFilter => ({
-  propertyName: filter.property_name,
-  operator: filter.operator,
-  value: filter.value,
-});
+const mapPropertyFilters = (filters: ApiPropertyFilter[]): PropertyFilter[] =>
+  filters.map((filter) => ({
+    propertyName: filter.property_name,
+    operator: filter.operator,
+    value: filter.value,
+  }));
 
 const mapAggregations = (
   aggregations: ApiAggregationResult[][]
 ): AggregationResult[][] =>
-  aggregations.map((aggGroup) =>
-    aggGroup.map((result) => ({
+  aggregations.map((aggregationGroup) =>
+    aggregationGroup.map((result) => ({
       collection: result.collection,
       searchQuery: result.search_query,
       groupbyProperty: result.groupby_property,
       aggregations: result.aggregations.map(mapPropertyAggregation),
-      filters: result.filters.map(mapPropertyFilter),
+      filters: mapPropertyFilters(result.filters),
     }))
   );
 

src/query/response/response.ts

@@ -18,7 +18,7 @@ export type QueryAgentResponse = {
 export type SearchResult = {
   collection: string;
   queries: string[];
-  filters: PropertyFilter[];
+  filters: PropertyFilter[][];
   filterOperators: "AND" | "OR";
 };