diff --git a/src/client/api/index.ts b/src/client/api/index.ts index f544c28..0f753aa 100644 --- a/src/client/api/index.ts +++ b/src/client/api/index.ts @@ -1,2 +1,3 @@ export { resend } from "./email"; +export { search } from "./search"; export { upstash, openai, anthropic, custom } from "./llm"; diff --git a/src/client/api/search.test.ts b/src/client/api/search.test.ts new file mode 100644 index 0000000..8189bc4 --- /dev/null +++ b/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", + }, + }, + }); + }); +}); diff --git a/src/client/api/search.ts b/src/client/api/search.ts new file mode 100644 index 0000000..bab3ab4 --- /dev/null +++ b/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 { + 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); +}; diff --git a/src/client/api/types.ts b/src/client/api/types.ts index 3799024..68057dc 100644 --- a/src/client/api/types.ts +++ b/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> = { 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>>; diff --git a/src/client/api/utils.ts b/src/client/api/utils.ts index 01841b5..305a080 100644 --- a/src/client/api/utils.ts +++ b/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; diff --git a/src/client/client.ts b/src/client/client.ts index 525a2ec..0611a2a 100644 --- a/src/client/client.ts +++ b/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 = { 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;