feat: integrate search upsert-data

src/client/api/index.ts

@@ -1,2 +1,3 @@
 export { resend } from "./email";
+export { search } from "./search";
 export { upstash, openai, anthropic, custom } from "./llm";

src/client/api/search.test.ts

@@ -0,0 +1,118 @@
+import { describe, test } from "bun:test";
+import { Client } from "../client";
+import { MOCK_QSTASH_SERVER_URL, mockQStashServer } from "../workflow/test-utils";
+import { nanoid } from "../utils";
+import { search } from "./search";
+
+describe("search", () => {
+  const qstashToken = nanoid();
+  const searchToken = nanoid();
+  const apiUrl = "https://mock-search.upstash.io";
+  const indexName = "movies";
+
+  const globalHeader = "global-header";
+  const globalHeaderOverwritten = "global-header-overwritten";
+  const requestHeader = "request-header";
+
+  const globalHeaderValue = nanoid();
+  const overWrittenOldValue = nanoid();
+  const overWrittenNewValue = nanoid();
+  const requestHeaderValue = nanoid();
+
+  const client = new Client({
+    baseUrl: MOCK_QSTASH_SERVER_URL,
+    token: qstashToken,
+    headers: {
+      [globalHeader]: globalHeaderValue,
+      [globalHeaderOverwritten]: overWrittenOldValue,
+    },
+  });
+
+  test("should use search upsert", async () => {
+    await mockQStashServer({
+      execute: async () => {
+        await client.publishJSON({
+          api: {
+            name: "search",
+            provider: search({ apiUrl, token: searchToken, indexName }),
+          },
+          body: [
+            {
+              id: "movie-1",
+              content: {
+                title: "Inception",
+                description: "A thriller about dreams within dreams.",
+              },
+              metadata: { genre: "sci-fi", year: 2010 },
+            },
+            {
+              id: "movie-2",
+              content: {
+                title: "The Godfather",
+                description: "A story about a powerful Italian-American crime family.",
+              },
+              metadata: { genre: "crime", year: 1972 },
+            },
+            {
+              id: "movie-3",
+              content: {
+                title: "The Dark Knight",
+                description: "A tale of Batman's fight against the Joker.",
+              },
+              metadata: { genre: "action", year: 2008 },
+            },
+          ],
+          headers: {
+            "content-type": "application/json",
+            [globalHeaderOverwritten]: overWrittenNewValue,
+            [requestHeader]: requestHeaderValue,
+          },
+        });
+      },
+      responseFields: {
+        body: { messageId: "msgId" },
+        status: 200,
+      },
+      receivesRequest: {
+        method: "POST",
+        token: qstashToken,
+        url: `http://localhost:8080/v2/publish/${apiUrl}/upsert-data`,
+        body: [
+          {
+            id: "movie-1",
+            content: {
+              title: "Inception",
+              description: "A thriller about dreams within dreams.",
+            },
+            metadata: { genre: "sci-fi", year: 2010 },
+          },
+          {
+            id: "movie-2",
+            content: {
+              title: "The Godfather",
+              description: "A story about a powerful Italian-American crime family.",
+            },
+            metadata: { genre: "crime", year: 1972 },
+          },
+          {
+            id: "movie-3",
+            content: {
+              title: "The Dark Knight",
+              description: "A tale of Batman's fight against the Joker.",
+            },
+            metadata: { genre: "action", year: 2008 },
+          },
+        ],
+        headers: {
+          "content-type": "application/json",
+          authorization: `Bearer ${qstashToken}`,
+          "upstash-forward-authorization": `Bearer ${searchToken}`,
+          [`upstash-forward-${requestHeader}`]: requestHeaderValue,
+          [`upstash-forward-${globalHeader}`]: globalHeaderValue,
+          [`upstash-forward-${globalHeaderOverwritten}`]: overWrittenNewValue,
+          "upstash-method": "POST",
+        },
+      },
+    });
+  });
+});

src/client/api/search.ts

@@ -0,0 +1,39 @@
+import { BaseProvider } from "./base";
+import type { ProviderInfo, SearchOwner } from "./types";
+
+export class SearchProvider extends BaseProvider<"search", SearchOwner> {
+  public readonly apiKind = "search";
+  public readonly indexName: string;
+  public readonly method = "POST";
+
+  constructor(baseUrl: string, token: string, owner: SearchOwner, indexName: string) {
+    super(baseUrl, token, owner);
+    this.indexName = indexName;
+  }
+
+  getRoute(): string[] {
+    return ["upsert-data", this.indexName];
+  }
+
+  getHeaders(_options: unknown): Record<string, string> {
+    return {
+      authorization: `Bearer ${this.token}`,
+    };
+  }
+
+  onFinish(providerInfo: ProviderInfo, _options: unknown): ProviderInfo {
+    return providerInfo;
+  }
+}
+
+export const search = ({
+  apiUrl,
+  token,
+  indexName,
+}: {
+  apiUrl: string;
+  token: string;
+  indexName: string;
+}): SearchProvider => {
+  return new SearchProvider(apiUrl, token, "search", indexName);
+};

src/client/api/types.ts

@@ -28,8 +28,8 @@ export type ProviderInfo = {
   method: HTTPMethods;
 };
 
-export type ApiKind = "llm" | "email";
-export type Owner = EmailOwner | LLMOwner;
+export type ApiKind = "llm" | "email" | "search";
+export type Owner = EmailOwner | LLMOwner | SearchOwner;
 
 type PublishApi<TName extends ApiKind, TProvider extends BaseProvider<TName>> = {
   name: TName;
@@ -50,3 +50,9 @@ export type LLMOptions = {
   analytics?: { name: "helicone"; token: string };
 };
 export type PublishLLMApi = PublishApi<"llm", BaseProvider<"llm", LLMOwner>> & LLMOptions;
+
+/**
+ * Search
+ */
+export type SearchOwner = "search";
+export type PublishSearchApi = PublishApi<"search", BaseProvider<"search", SearchOwner>>;

src/client/api/utils.ts

@@ -1,5 +1,11 @@
 import type { PublishRequest } from "../client";
-import type { LLMOptions, ProviderInfo, PublishEmailApi, PublishLLMApi } from "./types";
+import type {
+  LLMOptions,
+  ProviderInfo,
+  PublishEmailApi,
+  PublishLLMApi,
+  PublishSearchApi,
+} from "./types";
 import { upstash } from "./llm";
 import type { HeadersInit } from "../types";
 
@@ -11,7 +17,7 @@ import type { HeadersInit } from "../types";
  * @returns updated request
  */
 export const getProviderInfo = (
-  api: PublishEmailApi | PublishLLMApi,
+  api: PublishEmailApi | PublishLLMApi | PublishSearchApi,
   upstashToken: string
 ): ProviderInfo => {
   const { name, provider, ...parameters } = api;

src/client/client.ts

@@ -23,7 +23,7 @@ import {
   wrapWithGlobalHeaders,
 } from "./utils";
 import { Workflow } from "./workflow";
-import type { PublishEmailApi, PublishLLMApi } from "./api/types";
+import type { PublishEmailApi, PublishLLMApi, PublishSearchApi } from "./api/types";
 import { processApi } from "./api/utils";
 import { VERSION } from "../../version";
 
@@ -279,6 +279,16 @@ export type PublishRequest<TBody = BodyInit> = {
       topic?: never;
       callback?: string;
     }
+  | {
+      url?: never;
+      urlGroup?: never;
+      /**
+       * The api endpoint the request should be sent to.
+       */
+      api: PublishSearchApi;
+      topic?: never;
+      callback?: string;
+    }
   | {
       url?: never;
       urlGroup?: never;