nosible 0.1.5

package/src/search/resultSet.ts

@@ -0,0 +1,387 @@
+import lunr from "lunr";
+import {type SearchQuery, type SearchResponse} from "../api/schemas";
+import {NosibleClient} from "../client";
+import {Result} from "./result";
+import {type FlattenResult} from "./result.types";
+import {AnalyzeBy, analyzeResults, type AnalyzeResult} from "./analyze";
+import {isBrowser} from "../utils";
+import {
+  exportJson,
+  exportNdjson,
+  exportCsv,
+  importJson,
+  importCsv,
+} from "../utils/file";
+import type {IResultSet} from "./resultSet.types";
+
+/**
+ * A collection of search results with methods for filtering, analysis, and export.
+ *
+ * ResultSet provides an interface for working with multiple search results,
+ * including full-text search, data analysis, and export to various formats.
+ *
+ * @example
+ * ```typescript
+ * // Create from search results
+ * const resultSet = ResultSet.fromResults(results, searchQuery, client);
+ *
+ * // Search within results
+ * const filtered = resultSet.findInResults('specific term');
+ *
+ * // Analyze results
+ * const analysis = resultSet.analyze(AnalyzeBy.SENTIMENT);
+ *
+ * // Export to CSV
+ * await resultSet.writeToCsv('results.csv');
+ * ```
+ */
+export class ResultSet implements IResultSet<Result> {
+  /** Array of Result objects in this set */
+  public results: Result[];
+
+  /**
+   * Creates a new ResultSet instance
+   *
+   * @param results - Array of Result objects to include in the set
+   */
+  constructor(results: Result[]) {
+    this.results = results;
+  }
+
+  /**
+   * Creates a ResultSet from a file path
+   *
+   * Supports JSON and CSV file formats. The file is read and parsed into Result objects.
+   *
+   * @param filePath - Path to the file to import
+   * @param client - NosibleClient instance for creating Result objects
+   * @returns Promise resolving to a new ResultSet
+   *
+   * @throws Error if file type is unsupported
+   *
+   * @example
+   * ```typescript
+   * const resultSet = await ResultSet.fromFilePath('./data.json', client);
+   * ```
+   */
+  static async fromFilePath(
+    filePath: string,
+    client: NosibleClient
+  ): Promise<ResultSet> {
+    // Handle JSON files
+    if (filePath.endsWith(".json")) {
+      const json = await importJson({filePath});
+      const results = json.map((result: any) => {
+        return new Result({client, result});
+      });
+      return new ResultSet(results);
+    } else if (filePath.endsWith(".csv")) {
+      const csv = await importCsv({filePath});
+      const results = csv.map((result: FlattenResult) => {
+        return Result.fromFlatten(client, result);
+      });
+      return new ResultSet(results);
+    } else {
+      throw new Error("Unsupported file type");
+    }
+  }
+
+  /**
+   * Creates a ResultSet from a File object (browser environment)
+   *
+   * Supports JSON and CSV file formats. The file is read and parsed into Result objects.
+   *
+   * @param file - File object to import
+   * @param client - NosibleClient instance for creating Result objects
+   * @returns Promise resolving to a new ResultSet
+   *
+   * @throws Error if file type is unsupported
+   *
+   * @example
+   * ```typescript
+   * const fileInput = document.getElementById('file') as HTMLInputElement;
+   * const file = fileInput.files[0];
+   * const resultSet = await ResultSet.fromFile(file, client);
+   * ```
+   */
+  static async fromFile(file: File, client: NosibleClient): Promise<ResultSet> {
+    if (file.name.endsWith(".json")) {
+      const json = await importJson({file});
+      const results = json.map((result: any) => {
+        return new Result({client, result});
+      });
+      return new ResultSet(results);
+    } else if (file.name.endsWith(".csv")) {
+      const csv = await importCsv({file});
+      const results = csv.map((result: any) => {
+        return new Result({client, result});
+      });
+      return new ResultSet(results);
+    } else {
+      throw new Error("Unsupported file type");
+    }
+  }
+
+  /**
+   * Creates a ResultSet from search results
+   *
+   * Converts SearchResponse objects and/or existing Result objects into a unified ResultSet.
+   *
+   * @param results - Array of SearchResponse or Result objects
+   * @param searchQuery - The original search query used (needed for SearchResponse objects)
+   * @param client - NosibleClient instance for creating Result objects
+   * @returns New ResultSet containing the processed results
+   *
+   * @example
+   * ```typescript
+   * const resultSet = ResultSet.fromResults(searchResults, query, client);
+   * ```
+   */
+  static fromResults(
+    results: (SearchResponse | Result)[],
+    searchQuery: SearchQuery,
+    client: NosibleClient
+  ): ResultSet {
+    const processedResults = results.map((result) => {
+      // Handle both SearchResponse and Result types
+      if (result instanceof Result) {
+        // If it's already a Result, return as is
+        return result;
+      } else {
+        // It's a SearchResponse, convert to Result
+        const fullResult = {...searchQuery, ...result};
+        return new Result({client, result: fullResult});
+      }
+    });
+    return new ResultSet(processedResults);
+  }
+
+  /**
+   * Combines multiple ResultSet objects into a single ResultSet
+   *
+   * @param resultSets - Array of ResultSet objects to merge
+   * @returns New ResultSet containing all results from the input sets
+   *
+   * @example
+   * ```typescript
+   * const combined = ResultSet.join([resultSet1, resultSet2, resultSet3]);
+   * ```
+   */
+  static join(resultSets: ResultSet[]) {
+    const newSet: Result[] = [];
+    resultSets.forEach((resultSet) => newSet.push(...resultSet.results));
+    return new ResultSet(newSet);
+  }
+
+  /**
+   * Returns the raw data objects from all results
+   *
+   * @returns Array of result data objects
+   */
+  public getResults = () => {
+    return this.results.map((result) => result.data());
+  };
+
+  /**
+   * Returns flattened results with all nested data expanded
+   *
+   * @returns Array of flattened result objects suitable for export
+   */
+  public getFlattenResults = (): FlattenResult[] => {
+    return this.results.map((result) => result.flatten());
+  };
+
+  /**
+   * Performs full-text search within the results using Lunr.js
+   *
+   * Creates a search index from result content and returns matching results.
+   * Useful for filtering large result sets based on specific terms.
+   *
+   * @param query - Search query string
+   * @returns Array of Result objects matching the query
+   *
+   * @example
+   * ```typescript
+   * const importantResults = resultSet.findInResults('urgent OR critical');
+   * ```
+   */
+  public findInResults = (query: string): Result[] => {
+    // Create index
+    const results = this.results;
+    const index = lunr(function (this: lunr.Builder) {
+      this.ref("index");
+      this.field("content");
+
+      results.forEach((result: Result, i: number) => {
+        this.add({
+          index: i,
+          content: result.content,
+        });
+      });
+    });
+
+    // Search index
+    const hits = index.search(query);
+    console.log("Number of results: " + this.results.length);
+
+    // match search results to original results
+    const matchedResults: Result[] = [];
+    hits.forEach((hit: lunr.Index.Result) => {
+      let index = Number(hit.ref);
+      const found = this.results[index];
+      if (found) {
+        matchedResults.push(found);
+      } else {
+        console.warn(`findInSearchResults: Hit not found (index ${index})`);
+      }
+    });
+    console.log("Number of hits: " + hits.length);
+    console.log("Number of matchedResults: " + matchedResults.length);
+
+    console.log("matchedResults: ", typeof matchedResults);
+    return matchedResults;
+  };
+
+  /**
+   * Analyzes results using various metrics and groupings
+   *
+   * @param by - Analysis type (sentiment, entities, topics, etc.)
+   * @returns Analysis result containing insights and statistics
+   *
+   * @example
+   * ```typescript
+   * const sentimentAnalysis = resultSet.analyze(AnalyzeBy.SENTIMENT);
+   * const entityAnalysis = resultSet.analyze(AnalyzeBy.ENTITIES);
+   * ```
+   */
+  public analyze = (by: AnalyzeBy): AnalyzeResult => {
+    return analyzeResults(this.results, by);
+  };
+
+  /**
+   * Converts results to a Polars DataFrame (Node.js only)
+   *
+   * Uses flattened results for tabular data structure. Not available in browser environments.
+   *
+   * @returns Promise resolving to a Polars DataFrame
+   *
+   * @throws Error if called in browser environment
+   *
+   * @example
+   * ```typescript
+   * const df = await resultSet.toPolars();
+   * const filtered = df.filter(pl.col('sentiment').gt(0.5));
+   * ```
+   */
+  public toPolars = async () => {
+    if (isBrowser()) {
+      throw new Error("toPolars is not supported in the browser");
+    } else {
+      const pl = await import("nodejs-polars");
+      const flattenedData = this.getFlattenResults();
+      const df = pl.DataFrame(flattenedData);
+      return df;
+    }
+  };
+
+  /**
+   * Exports results to a CSV file
+   *
+   * Uses flattened results for tabular format. Works in both Node.js and browser environments.
+   *
+   * @param fileName - Name of the CSV file to create
+   *
+   * @example
+   * ```typescript
+   * await resultSet.writeToCsv('search-results.csv');
+   * ```
+   */
+  public writeToCsv = async (fileName: string) => {
+    const flattenedData = this.getFlattenResults();
+    await exportCsv(flattenedData, fileName);
+  };
+
+  /**
+   * Exports results to a Parquet file (Node.js only)
+   *
+   * Creates an efficient columnar storage format using Polars. Not available in browser environments.
+   *
+   * @param fileName - Name of the Parquet file to create
+   *
+   * @throws Error if called in browser environment
+   *
+   * @example
+   * ```typescript
+   * await resultSet.writeToParquet('results.parquet');
+   * ```
+   */
+  public writeToParquet = async (fileName: string) => {
+    if (isBrowser()) {
+      throw new Error("writeToParquet is not supported in the browser");
+    } else {
+      const pl = await import("nodejs-polars");
+      const flattenedData = this.getFlattenResults();
+      const df = pl.DataFrame(flattenedData);
+      await df.writeParquet(fileName);
+    }
+  };
+
+  /**
+   * Exports results to an IPC file (Node.js only)
+   *
+   * Creates an Arrow IPC format file using Polars. Not available in browser environments.
+   *
+   * @param fileName - Name of the IPC file to create
+   *
+   * @throws Error if called in browser environment
+   *
+   * @example
+   * ```typescript
+   * await resultSet.writeToIpc('results.arrow');
+   * ```
+   */
+  public writeToIpc = async (fileName: string) => {
+    if (isBrowser()) {
+      throw new Error("writeToIpc is not supported in the browser");
+    } else {
+      const pl = await import("nodejs-polars");
+      const flattenedData = this.getFlattenResults();
+      const df = pl.DataFrame(flattenedData);
+      await df.writeIPC(fileName);
+    }
+  };
+
+  /**
+   * Exports results to a JSON file
+   *
+   * Uses the raw result data format. Works in both Node.js and browser environments.
+   *
+   * @param fileName - Name of the JSON file to create
+   *
+   * @example
+   * ```typescript
+   * await resultSet.writeToJson('results.json');
+   * ```
+   */
+  public writeToJson = async (fileName: string) => {
+    await exportJson(this.getResults(), fileName);
+  };
+
+  /**
+   * Exports results to a newline-delimited JSON (NDJSON) file
+   *
+   * Uses flattened results with each result on a separate line.
+   * Works in both Node.js and browser environments.
+   *
+   * @param fileName - Name of the NDJSON file to create
+   *
+   * @example
+   * ```typescript
+   * await resultSet.writeToNdjson('results.ndjson');
+   * ```
+   */
+  public writeToNdjson = async (fileName: string) => {
+    const flattenedData = this.getFlattenResults();
+    await exportNdjson(flattenedData, fileName);
+  };
+}

package/src/search/resultSet.types.ts

@@ -0,0 +1,3 @@
+export interface IResultSet<T = any> {
+  results: T[];
+}

package/src/search/search.test.ts

@@ -0,0 +1,299 @@
+import {describe, it, expect, beforeEach} from "bun:test";
+import {Search, type UserSearchParamsType} from "./search";
+import type {NosibleClient} from "../client";
+import {mockSearchQuery} from "../test-utils/mocks";
+
+// Re-export for backward compatibility
+export {mockSearchQuery};
+
+describe("Search", () => {
+  let mockClient: NosibleClient;
+  let searchParams: UserSearchParamsType;
+
+  beforeEach(() => {
+    // Create a mock client with minimal required properties
+    mockClient = {
+      expansionsModel: "openai/gpt-4o",
+      llmClient: {
+        chat: {
+          completions: {
+            create: async () => ({
+              choices: [{message: {content: '["test expansion"]'}}],
+            }),
+          },
+        },
+      },
+    } as unknown as NosibleClient;
+
+    // Default search parameters
+    searchParams = {
+      question: "artificial intelligence machine learning",
+      nResults: 10,
+    };
+  });
+
+  describe("constructor", () => {
+    it("should create a Search instance with basic parameters", () => {
+      const search = new Search(searchParams);
+
+      expect(search.question).toBe("artificial intelligence machine learning");
+      expect(search.nResults).toBe(10);
+    });
+
+    it("should create a Search instance with all optional parameters", () => {
+      const fullParams: UserSearchParamsType = {
+        ...searchParams,
+        autoGenerateExpansions: true,
+        expansions: ["AI", "ML", "neural networks"],
+        nProbes: 5,
+        nContextify: 256,
+        algorithm: "hybrid-2",
+        sqlFilter: "published >= '2023-01-01'",
+        minSimilarity: 0.7,
+        mustInclude: ["deep learning", "neural networks"],
+        mustExclude: ["basic", "beginner"],
+        brandSafety: "safe",
+        language: "en",
+        continent: "North America",
+        region: "North America",
+        country: "US",
+        sector: "Technology",
+        industryGroup: "Software",
+        industry: "Artificial Intelligence",
+        subIndustry: "Machine Learning",
+        iabTier1: "Technology",
+        iabTier2: "Computing",
+        iabTier3: "Artificial Intelligence",
+        iabTier4: "Machine Learning",
+        companies: ["GOOG", "MSFT"],
+        publishStart: new Date("2023-01-01"),
+        publishEnd: new Date("2023-12-31"),
+        visitedStart: new Date("2023-01-01"),
+        visitedEnd: new Date("2023-12-31"),
+        certain: true,
+        includeNetlocs: ["arxiv.org", "ieee.org"],
+        excludeNetlocs: ["spam.com"],
+        includeCompanies: ["GOOG", "MSFT"],
+        excludeCompanies: ["SPAM"],
+        includeDocs: ["hash1", "hash2"],
+        excludeDocs: ["hash3"],
+      };
+
+      const search = new Search(fullParams);
+
+      expect(search.question).toBe(fullParams.question as string);
+      expect(search.nResults).toBe(fullParams.nResults as number);
+      expect(search.autoGenerateExpansions).toBe(true);
+      expect(search.expansions).toEqual(fullParams.expansions);
+      expect(search.nProbes).toBe(5);
+      expect(search.nContextify).toBe(256);
+      expect(search.algorithm).toBe("hybrid-2");
+      expect(search.sqlFilter).toBe(fullParams.sqlFilter);
+      expect(search.minSimilarity).toBe(0.7);
+      expect(search.mustInclude).toEqual(fullParams.mustInclude);
+      expect(search.mustExclude).toEqual(fullParams.mustExclude);
+      expect(search.brandSafety).toBe("safe");
+      expect(search.language).toBe("en");
+      expect(search.continent).toBe("North America");
+      expect(search.region).toBe("North America");
+      expect(search.country).toBe("US");
+      expect(search.sector).toBe("Technology");
+      expect(search.industryGroup).toBe("Software");
+      expect(search.industry).toBe("Artificial Intelligence");
+      expect(search.subIndustry).toBe("Machine Learning");
+      expect(search.iabTier1).toBe("Technology");
+      expect(search.iabTier2).toBe("Computing");
+      expect(search.iabTier3).toBe("Artificial Intelligence");
+      expect(search.iabTier4).toBe("Machine Learning");
+      expect(search.companies).toEqual(fullParams.companies);
+      expect(search.publishStart).toEqual(fullParams.publishStart);
+      expect(search.publishEnd).toEqual(fullParams.publishEnd);
+      expect(search.visitedStart).toEqual(fullParams.visitedStart);
+      expect(search.visitedEnd).toEqual(fullParams.visitedEnd);
+      expect(search.certain).toBe(true);
+      expect(search.includeNetlocs).toEqual(fullParams.includeNetlocs);
+      expect(search.excludeNetlocs).toEqual(fullParams.excludeNetlocs);
+      expect(search.includeCompanies).toEqual(fullParams.includeCompanies);
+      expect(search.excludeCompanies).toEqual(fullParams.excludeCompanies);
+      expect(search.includeDocs).toEqual(fullParams.includeDocs);
+      expect(search.excludeDocs).toEqual(fullParams.excludeDocs);
+    });
+  });
+
+  describe("searchBody", () => {
+    it("should generate search body with basic parameters", async () => {
+      const search = new Search(searchParams);
+
+      const body = await search.searchBody({client: mockClient});
+
+      expect(body).toEqual({
+        question: "artificial intelligence machine learning",
+        expansions: undefined,
+        n_results: 10,
+      });
+    });
+
+    it("should generate search body with predefined expansions", async () => {
+      const search = new Search({
+        ...searchParams,
+        expansions: ["AI", "ML"],
+      });
+
+      const body = await search.searchBody({client: mockClient});
+
+      expect(body).toEqual({
+        question: "artificial intelligence machine learning",
+        expansions: ["AI", "ML"],
+        n_results: 10,
+      });
+    });
+
+    it("should generate search body with all parameters", async () => {
+      const search = new Search({
+        ...searchParams,
+        nProbes: 5,
+        nContextify: 256,
+        algorithm: "hybrid-2",
+        minSimilarity: 0.7,
+        mustInclude: ["deep learning"],
+        mustExclude: ["basic"],
+        brandSafety: "safe",
+        language: "en",
+        continent: "North America",
+        region: "North America",
+        country: "US",
+        sector: "Technology",
+        industryGroup: "Software",
+        industry: "AI",
+        subIndustry: "ML",
+        iabTier1: "Technology",
+        iabTier2: "Computing",
+        iabTier3: "AI",
+        iabTier4: "ML",
+        companies: ["GOOG"],
+        publishStart: new Date("2023-01-01"),
+        publishEnd: new Date("2023-12-31"),
+        visitedStart: new Date("2023-01-01"),
+        visitedEnd: new Date("2023-12-31"),
+        certain: true,
+        includeNetlocs: ["arxiv.org"],
+        excludeNetlocs: ["spam.com"],
+        includeCompanies: ["GOOG"],
+        excludeCompanies: ["SPAM"],
+        includeDocs: ["hash1"],
+        excludeDocs: ["hash2"],
+      });
+
+      const body = await search.searchBody({client: mockClient});
+
+      expect(body).toEqual({
+        question: "artificial intelligence machine learning",
+        expansions: undefined,
+        sql_filter:
+          "SELECT loc FROM engine WHERE published >= '2023-01-01T00:00:00.000Z' AND published <= '2023-12-31T00:00:00.000Z' AND visited >= '2023-01-01T00:00:00.000Z' AND visited <= '2023-12-31T00:00:00.000Z' AND certain = TRUE AND netloc IN ('arxiv.org', 'www.arxiv.org') AND netloc NOT IN ('spam.com', 'www.spam.com') AND (companies IS NOT NULL AND (ARRAY_CONTAINS(companies, 'GOOG'))) AND (companies IS NULL OR NOT (ARRAY_CONTAINS(companies, 'SPAM'))) AND doc_hash IN ('hash1') AND doc_hash NOT IN ('hash2')",
+        algorithm: "hybrid-2",
+        n_probes: 5,
+        n_results: 10,
+        n_contextify: 256,
+        min_similarity: 0.7,
+        must_include: ["deep learning"],
+        must_exclude: ["basic"],
+        brand_safety: "safe",
+        language: "en",
+        continent: "North America",
+        region: "North America",
+        country: "US",
+        sector: "Technology",
+        industry_group: "Software",
+        industry: "AI",
+        sub_industry: "ML",
+        iab_tier_1: "Technology",
+        iab_tier_2: "Computing",
+        iab_tier_3: "AI",
+        iab_tier_4: "ML",
+        companies: ["GOOG"],
+      });
+    });
+
+    it("should filter out undefined values", async () => {
+      const search = new Search(searchParams);
+
+      const body = await search.searchBody({client: mockClient});
+
+      // Check that undefined values are not present in the output
+      expect(body).not.toHaveProperty("expansions");
+      expect(body).not.toHaveProperty("algorithm");
+      expect(body).not.toHaveProperty("n_probes");
+      expect(body).not.toHaveProperty("n_contextify");
+      expect(body).not.toHaveProperty("min_similarity");
+      expect(body).not.toHaveProperty("must_include");
+      expect(body).not.toHaveProperty("must_exclude");
+      expect(body).not.toHaveProperty("brand_safety");
+      expect(body).not.toHaveProperty("language");
+      expect(body).not.toHaveProperty("continent");
+      expect(body).not.toHaveProperty("region");
+      expect(body).not.toHaveProperty("country");
+      expect(body).not.toHaveProperty("sector");
+      expect(body).not.toHaveProperty("industry_group");
+      expect(body).not.toHaveProperty("industry");
+      expect(body).not.toHaveProperty("sub_industry");
+      expect(body).not.toHaveProperty("iab_tier_1");
+      expect(body).not.toHaveProperty("iab_tier_2");
+      expect(body).not.toHaveProperty("iab_tier_3");
+      expect(body).not.toHaveProperty("iab_tier_4");
+      expect(body).not.toHaveProperty("companies");
+
+      // But should have required fields
+      expect(body).toHaveProperty("question");
+      expect(body).toHaveProperty("n_results");
+    });
+
+    it("should handle boundary values for numeric parameters", async () => {
+      const search = new Search({
+        ...searchParams,
+        nProbes: 1, // min value
+        nContextify: 64, // min value
+        minSimilarity: 0, // min value
+      });
+
+      const body = await search.searchBody({client: mockClient});
+
+      expect(body.n_probes).toBe(1);
+      expect(body.n_contextify).toBe(64);
+      expect(body.min_similarity).toBe(0);
+    });
+
+    it("should handle maximum values for numeric parameters", async () => {
+      const search = new Search({
+        ...searchParams,
+        nProbes: 10, // max value
+        nContextify: 1024, // max value
+        minSimilarity: 1, // max value
+      });
+
+      const body = await search.searchBody({client: mockClient});
+
+      expect(body.n_probes).toBe(10);
+      expect(body.n_contextify).toBe(1024);
+      expect(body.min_similarity).toBe(1);
+    });
+
+    it("should handle autoGenerateExpansions without LLM client", async () => {
+      const search = new Search({
+        ...searchParams,
+        autoGenerateExpansions: true,
+      });
+
+      const clientWithoutLLM = {
+        expansionsModel: "openai/gpt-4o",
+        llmClient: undefined,
+      } as unknown as NosibleClient;
+
+      await expect(
+        search.searchBody({client: clientWithoutLLM})
+      ).rejects.toThrow(
+        "A OpenAI client instance must be provided as 'llmClient'."
+      );
+    });
+  });
+});