nosible 0.1.5

package/src/search/search.ts

@@ -0,0 +1,187 @@
+import {generateSqlFilter, sqlFilterSchema} from "./sqlFilter";
+import {z} from "zod";
+import {
+  algorithmEnum,
+  brandSafetyEnum,
+  continentEnum,
+  regionEnum,
+  type AlgorithmEnum,
+  type BrandSafetyEnum,
+  type ContinentEnum,
+  type RegionEnum,
+} from "../api/schemas";
+import type {NosibleClient} from "../client";
+import {generateExpansions} from "../utils/llm";
+import {aiSearchParamsSchema, type StdSearchParamsType} from "../api/index";
+
+/**
+ * User Search Params
+ * A user can search using additional filters, these are passed as a SQL filter string.
+ * */
+// User AI Search Params
+export const userAiSearchParamsSchema = aiSearchParamsSchema;
+export type UserAiSearchParamsType = z.infer<typeof userAiSearchParamsSchema>;
+
+// User Fast Search Params
+export const userSearchParamsSchema = sqlFilterSchema.extend({
+  // Standard search params for fast and bulk search only
+  question: z.string().optional(),
+  autoGenerateExpansions: z.boolean().optional(),
+  expansions: z.array(z.string()).optional(),
+  sqlFilter: z.string().optional(),
+  algorithm: algorithmEnum.optional(),
+  nResults: z.number().optional(),
+  nProbes: z.number().min(1).max(10).optional(),
+  nContextify: z.number().min(64).max(1024).optional(),
+  minSimilarity: z.number().min(0).max(1).optional(),
+  mustInclude: z.array(z.string()).max(100).optional(),
+  mustExclude: z.array(z.string()).max(100).optional(),
+  brandSafety: brandSafetyEnum.optional(),
+  language: z.string().min(2).max(2).optional(),
+  continent: continentEnum.optional(),
+  region: regionEnum.optional(),
+  country: z.string().optional(),
+  sector: z.string().optional(),
+  industryGroup: z.string().optional(),
+  industry: z.string().optional(),
+  subIndustry: z.string().optional(),
+  iabTier1: z.string().optional(),
+  iabTier2: z.string().optional(),
+  iabTier3: z.string().optional(),
+  iabTier4: z.string().optional(),
+  companies: z.array(z.string()).max(30).optional(),
+});
+export type UserSearchParamsType = z.infer<typeof userSearchParamsSchema>;
+
+export const clientDefaultSearchParamsSchema = userSearchParamsSchema.omit({
+  question: true,
+});
+
+export type ClientDefaultSearchParamsType = z.infer<
+  typeof clientDefaultSearchParamsSchema
+>;
+
+export type ClientSearchParamsSchema =
+  | (UserSearchParamsType & {search?: never})
+  | {search: Search};
+
+// Similar Search User Params
+const similarUserSearchParams = userSearchParamsSchema.omit({
+  question: true,
+});
+
+export type SimilarUserSearchParamsType = z.infer<
+  typeof similarUserSearchParams
+>;
+
+export class Search {
+  question!: string;
+  nResults!: number;
+  expansions?: string[];
+  autoGenerateExpansions?: boolean;
+  nProbes?: number;
+  nContextify?: number;
+  algorithm?: AlgorithmEnum;
+  sqlFilter?: string;
+  minSimilarity?: number;
+  mustInclude?: string[];
+  mustExclude?: string[];
+  brandSafety?: BrandSafetyEnum;
+  language?: string;
+  continent?: ContinentEnum;
+  region?: RegionEnum;
+  country?: string;
+  sector?: string;
+  industryGroup?: string;
+  industry?: string;
+  subIndustry?: string;
+  iabTier1?: string;
+  iabTier2?: string;
+  iabTier3?: string;
+  iabTier4?: string;
+  companies?: string[];
+  // SQL Filter
+  publishStart?: Date;
+  publishEnd?: Date;
+  visitedStart?: Date;
+  visitedEnd?: Date;
+  certain?: boolean;
+  includeNetlocs?: string[];
+  excludeNetlocs?: string[];
+  includeCompanies?: string[];
+  excludeCompanies?: string[];
+  includeDocs?: string[];
+  excludeDocs?: string[];
+
+  constructor(params: UserSearchParamsType) {
+    Object.assign(this, params);
+  }
+
+  /* Returns the search body for the API call
+
+  Generates the expansions and sql filter from the search params if required
+   */
+  public searchBody = async ({
+    client,
+  }: {
+    client: NosibleClient;
+  }): Promise<StdSearchParamsType> => {
+    let expansions: string[] | undefined = this.expansions;
+
+    if (this.autoGenerateExpansions) {
+      expansions = await generateExpansions({
+        question: this.question,
+        model: client.expansionsModel,
+        llmClient: client.llmClient,
+      });
+    }
+
+    let sqlFilter = generateSqlFilter({
+      publishStart: this.publishStart || undefined,
+      publishEnd: this.publishEnd || undefined,
+      visitedStart: this.visitedStart || undefined,
+      visitedEnd: this.visitedEnd || undefined,
+      certain: this.certain || undefined,
+      includeNetlocs: this.includeNetlocs || undefined,
+      excludeNetlocs: this.excludeNetlocs || undefined,
+      includeCompanies: this.includeCompanies || undefined,
+      excludeCompanies: this.excludeCompanies || undefined,
+      includeDocs: this.includeDocs || undefined,
+      excludeDocs: this.excludeDocs || undefined,
+    });
+
+    const output: StdSearchParamsType = {
+      question: this.question,
+      expansions: expansions,
+      sql_filter: sqlFilter || undefined,
+      algorithm: this.algorithm,
+      n_probes: this.nProbes,
+      n_results: this.nResults,
+      n_contextify: this.nContextify,
+      min_similarity: this.minSimilarity,
+      must_include: this.mustInclude,
+      must_exclude: this.mustExclude,
+      brand_safety: this.brandSafety,
+      language: this.language,
+      continent: this.continent,
+      region: this.region,
+      country: this.country,
+      sector: this.sector,
+      industry_group: this.industryGroup,
+      industry: this.industry,
+      sub_industry: this.subIndustry,
+      iab_tier_1: this.iabTier1,
+      iab_tier_2: this.iabTier2,
+      iab_tier_3: this.iabTier3,
+      iab_tier_4: this.iabTier4,
+      companies: this.companies,
+    };
+
+    // remove undefined values
+    const outputWithoutUndefined = Object.fromEntries(
+      Object.entries(output).filter(([, value]) => value !== undefined)
+    ) as StdSearchParamsType;
+
+    return outputWithoutUndefined;
+  };
+}

package/src/search/searchSet.io.test.ts

@@ -0,0 +1,321 @@
+import {describe, it, expect, beforeAll} from "bun:test";
+import {SearchSet} from "./searchSet";
+import {Search} from "./search";
+import {importCsv, importJson} from "../utils/file";
+import {randomUUID} from "node:crypto";
+
+/**
+ * Test configuration
+ */
+const TEST_DIR = `${process.cwd()}/tmp/searchSet/`; // Directory where test files will be created
+const CLEANUP = false; // Set to false to keep test files for inspection
+
+/**
+ * Helper function to conditionally cleanup test files
+ */
+const cleanupFile = async (filePath: string) => {
+  if (CLEANUP) {
+    const fs = await import("node:fs/promises");
+    await fs.unlink(filePath).catch(() => {});
+  } else {
+    console.log(`Keeping test file: ${filePath}`);
+  }
+};
+
+/**
+ * I/O tests for SearchSet with file system operations
+ * These tests require file system access for reading/writing
+ */
+describe("SearchSet - I/O Tests", () => {
+  let testSearchSet: SearchSet;
+
+  beforeAll(async () => {
+    try {
+      // Create test directory if it doesn't exist
+      const fs = await import("node:fs/promises");
+      await fs.mkdir(TEST_DIR, {recursive: true}).catch(() => {});
+
+      // Create a test SearchSet with sample searches
+      const searches = [
+        new Search({
+          question: "artificial intelligence",
+          nResults: 10,
+        }),
+        new Search({
+          question: "machine learning",
+          nResults: 20,
+          algorithm: "hybrid-1",
+        }),
+        new Search({
+          question: "deep learning",
+          nResults: 15,
+          language: "en",
+          brandSafety: "safe",
+        }),
+      ];
+
+      testSearchSet = new SearchSet(searches);
+    } catch (error) {
+      console.warn("Skipping I/O tests due to setup error");
+      throw error;
+    }
+  });
+
+  describe("File export methods", () => {
+    it("should export SearchSet to JSON file", async () => {
+      const uniqueId = randomUUID();
+      const fileName = `${TEST_DIR}/test-searches-${uniqueId}.json`;
+      await testSearchSet.writeToJson(fileName);
+
+      console.log(
+        `Successfully exported ${testSearchSet.searches.length} searches to JSON`
+      );
+
+      // Verify the file was created
+      const fs = await import("node:fs/promises");
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      const fileExists = await fs
+        .access(fileName)
+        .then(() => true)
+        .catch(() => false);
+      expect(fileExists).toBe(true);
+
+      // Clean up
+      await cleanupFile(fileName);
+    });
+
+    it("should export SearchSet to NDJSON file", async () => {
+      const uniqueId = randomUUID();
+      const fileName = `${TEST_DIR}/test-searches-${uniqueId}.ndjson`;
+      await testSearchSet.writeToNdjson(fileName);
+
+      console.log(
+        `Successfully exported ${testSearchSet.searches.length} searches to NDJSON`
+      );
+
+      // Verify the file was created
+      const fs = await import("node:fs/promises");
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      const fileExists = await fs
+        .access(fileName)
+        .then(() => true)
+        .catch(() => false);
+      expect(fileExists).toBe(true);
+
+      // Clean up
+      await cleanupFile(fileName);
+    });
+
+    it("should export SearchSet to CSV file", async () => {
+      const uniqueId = randomUUID();
+      const fileName = `${TEST_DIR}/test-searches-${uniqueId}.csv`;
+      await testSearchSet.writeToCsv(fileName);
+
+      console.log(
+        `Successfully exported ${testSearchSet.searches.length} searches to CSV`
+      );
+
+      // Verify the file was created
+      const fs = await import("node:fs/promises");
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      const fileExists = await fs
+        .access(fileName)
+        .then(() => true)
+        .catch(() => false);
+      expect(fileExists).toBe(true);
+
+      // Read file and verify contents
+      const json = await importCsv({filePath: fileName});
+      expect(json).toBeDefined();
+      expect(json.length).toBeGreaterThan(0);
+      expect(json[0]?.question).toBeDefined();
+
+      // Clean up
+      await cleanupFile(fileName);
+    });
+  });
+
+  describe("Data conversion methods", () => {
+    it("should convert SearchSet to Polars DataFrame", async () => {
+      const df = await testSearchSet.toPolars();
+
+      expect(df).toBeDefined();
+      expect(df.height).toBeGreaterThan(0);
+      expect(df.width).toBeGreaterThan(0);
+
+      console.log(
+        `Successfully converted ${testSearchSet.searches.length} searches to Polars DataFrame`
+      );
+      console.log(`DataFrame shape: ${df.height} rows x ${df.width} columns`);
+    });
+  });
+
+  describe("Round-trip export/import - JSON", () => {
+    it("should verify exported JSON can be reimported", async () => {
+      const uniqueId = randomUUID();
+      const fileName = `${TEST_DIR}/test-roundtrip-${uniqueId}.json`;
+      await testSearchSet.writeToJson(fileName);
+
+      // Import the JSON file back
+      const reimportedSearchSet = await SearchSet.fromFilePath(fileName);
+
+      expect(reimportedSearchSet).toBeDefined();
+      expect(reimportedSearchSet.searches.length).toBe(
+        testSearchSet.searches.length
+      );
+
+      console.log(
+        `Successfully round-tripped ${reimportedSearchSet.searches.length} searches through JSON`
+      );
+
+      // Clean up
+      await cleanupFile(fileName);
+    });
+
+    it("should produce identical JSON files after save -> load -> save cycle", async () => {
+      const fs = await import("node:fs/promises");
+      const uniqueId = randomUUID();
+      const file1 = `${TEST_DIR}/test-roundtrip-1-${uniqueId}.json`;
+      const file2 = `${TEST_DIR}/test-roundtrip-2-${uniqueId}.json`;
+
+      try {
+        // Save original searches
+        await testSearchSet.writeToJson(file1);
+
+        // Load and verify
+        const reloadedSearchSet = await SearchSet.fromFilePath(file1);
+        expect(reloadedSearchSet.searches.length).toBe(
+          testSearchSet.searches.length
+        );
+
+        // Save again
+        await reloadedSearchSet.writeToJson(file2);
+
+        // Compare the two JSON files
+        const json1Content = await fs.readFile(file1, "utf-8");
+        const json2Content = await fs.readFile(file2, "utf-8");
+
+        const json1 = JSON.parse(json1Content);
+        const json2 = JSON.parse(json2Content);
+
+        // Verify they have the same structure and content
+        expect(json1.length).toBe(json2.length);
+        expect(json1).toEqual(json2);
+
+        console.log(
+          `✓ Both JSON files are identical (${json1.length} searches)`
+        );
+      } finally {
+        // Clean up both files
+        await cleanupFile(file1);
+        await cleanupFile(file2);
+      }
+    });
+  });
+
+  describe("Round-trip export/import - CSV", () => {
+    it("should verify exported CSV can be reimported", async () => {
+      const uniqueId = randomUUID();
+      const fileName = `${TEST_DIR}/test-roundtrip-${uniqueId}.csv`;
+      await testSearchSet.writeToCsv(fileName);
+
+      // Import the CSV file back
+      const reimportedSearchSet = await SearchSet.fromFilePath(fileName);
+
+      expect(reimportedSearchSet).toBeDefined();
+      expect(reimportedSearchSet.searches.length).toBe(
+        testSearchSet.searches.length
+      );
+
+      console.log(
+        `Successfully round-tripped ${reimportedSearchSet.searches.length} searches through CSV`
+      );
+
+      // Clean up
+      await cleanupFile(fileName);
+    });
+
+    it("should produce identical CSV files after save -> load -> save cycle", async () => {
+      const fs = await import("node:fs/promises");
+      const uniqueId = randomUUID();
+      const file1 = `${TEST_DIR}/test-roundtrip-1-${uniqueId}.csv`;
+      const file2 = `${TEST_DIR}/test-roundtrip-2-${uniqueId}.csv`;
+
+      try {
+        // Save original searches
+        await testSearchSet.writeToCsv(file1);
+
+        // Load and verify
+        const reloadedSearchSet = await SearchSet.fromFilePath(file1);
+        expect(reloadedSearchSet.searches.length).toBe(
+          testSearchSet.searches.length
+        );
+
+        // Save again
+        await reloadedSearchSet.writeToCsv(file2);
+
+        // Compare the two CSV files
+        const csv1Content = await fs.readFile(file1, "utf-8");
+        const csv2Content = await fs.readFile(file2, "utf-8");
+
+        // Verify they have the same content
+        expect(csv1Content).toBe(csv2Content);
+
+        console.log(`✓ Both CSV files are identical`);
+      } finally {
+        // Clean up both files
+        await cleanupFile(file1);
+        await cleanupFile(file2);
+      }
+    });
+  });
+
+  describe("SearchSet with various search parameters", () => {
+    it("should handle searches with complex filter parameters", async () => {
+      const complexSearches = [
+        new Search({
+          question: "climate change",
+          nResults: 50,
+          algorithm: "hybrid-1",
+          brandSafety: "safe",
+          language: "en",
+          country: "US",
+          sector: "Technology",
+          mustInclude: ["renewable", "energy"],
+          mustExclude: ["fossil"],
+        }),
+        new Search({
+          question: "renewable energy",
+          nResults: 30,
+          expansions: ["solar power", "wind energy", "hydroelectric"],
+          minSimilarity: 0.7,
+        }),
+      ];
+
+      const complexSearchSet = new SearchSet(complexSearches);
+      const uniqueId = randomUUID();
+      const fileName = `${TEST_DIR}/complex-searches-${uniqueId}.json`;
+
+      await complexSearchSet.writeToJson(fileName);
+
+      // Verify round-trip
+      const reimported = await SearchSet.fromFilePath(fileName);
+      expect(reimported.searches.length).toBe(complexSearches.length);
+      expect(reimported.searches[0]?.question).toBe("climate change");
+      expect(reimported.searches[0]?.mustInclude).toEqual([
+        "renewable",
+        "energy",
+      ]);
+
+      console.log(
+        `✓ Successfully handled complex search parameters (${reimported.searches.length} searches)`
+      );
+
+      // Clean up
+      await cleanupFile(fileName);
+    });
+  });
+});

package/src/search/searchSet.ts

@@ -0,0 +1,122 @@
+import {
+  exportJson,
+  exportCsv,
+  exportNdjson,
+  importCsv,
+  importJson,
+} from "../utils/file";
+import {Search} from "./search";
+import {z} from "zod";
+import {isBrowser} from "../utils";
+
+/**
+ * A set of searches.
+ */
+
+const searchSetParams = z.union([
+  z.object({
+    searches: z.array(z.instanceof(Search)),
+  }),
+  z.array(z.instanceof(Search)),
+]);
+
+export type SearchSetParams = z.infer<typeof searchSetParams>;
+
+export class SearchSet {
+  public searches: Search[] = [];
+  public index: number;
+
+  constructor(params: SearchSetParams) {
+    // Handle both object and array forms
+    if (Array.isArray(params)) {
+      this.searches = params;
+    } else {
+      this.searches = params.searches;
+    }
+
+    this.index = 0;
+  }
+
+  /**
+   * Creates a SearchSet from a JSON file asynchronously
+   * @param filePath Path to the JSON file
+   * @returns Promise<SearchSet>
+   */
+  static async fromFilePath(filePath: string): Promise<SearchSet> {
+    if (filePath.endsWith(".json")) {
+      const searches = await importJson({filePath: filePath});
+      return new SearchSet(searches.map((search: any) => new Search(search)));
+    } else if (filePath.endsWith(".csv")) {
+      const searches = await importCsv({filePath: filePath});
+      return new SearchSet(searches.map((search: any) => new Search(search)));
+    } else throw new Error("Unsupported file type");
+  }
+
+  /**
+   * Creates a SearchSet from a File asynchronously
+   * @param file File to read from
+   * @returns Promise<SearchSet>
+   */
+
+  static async fromFile(file: File): Promise<SearchSet> {
+    if (file.name.endsWith(".json")) {
+      const searches = await importJson({file: file});
+      return new SearchSet(searches.map((search: any) => new Search(search)));
+    } else if (file.name.endsWith(".csv")) {
+      const searches = await importCsv({file: file});
+      return new SearchSet(searches.map((search: any) => new Search(search)));
+    } else throw new Error("Unsupported file type");
+  }
+
+  /**
+   * Returns the next search in the set.
+   * @returns The next search in the set, or undefined if there are no more searches.
+   */
+  public next = (): Search | null => {
+    this.index++;
+    const search = this.searches[this.index];
+    if (!search) {
+      return null;
+    }
+    return search;
+  };
+
+  /**
+   * Adds a search to the set.
+   * @param search The search to add.
+   */
+  public add = (search: Search) => {
+    this.searches.push(search);
+  };
+
+  /**
+   * Removes a search from the set.
+   * @param index The index of the search to remove.
+   */
+  public remove = (index: number) => {
+    this.searches.splice(index, 1);
+  };
+
+  public toPolars = async () => {
+    if (isBrowser()) {
+      throw new Error("toPolars is not supported in the browser");
+    } else {
+      const pl = await import("nodejs-polars");
+      const df = pl.DataFrame(this.searches);
+      return df;
+    }
+  };
+
+  // Export to file or download
+  public writeToCsv = async (fileName: string) => {
+    await exportCsv(this.searches, fileName);
+  };
+
+  public writeToJson = async (fileName: string) => {
+    await exportJson(this.searches, fileName);
+  };
+
+  public writeToNdjson = async (fileName: string) => {
+    await exportNdjson(this.searches, fileName);
+  };
+}