├── .eslintrc.js ├── .github └── workflows │ └── ci_test.yml ├── .gitignore ├── LICENSE ├── README.md ├── jest_jsdom.config.js ├── jest_node.config.js ├── package-lock.json ├── package.json ├── src ├── createHTTPClient.ts ├── helpers.ts ├── httpClient │ ├── default.ts │ └── node.ts ├── index.ts ├── turbopuffer.test.ts ├── turbopuffer.ts └── types.ts └── tsconfig.json /.eslintrc.js: -------------------------------------------------------------------------------- 1 | /** @type {import("eslint").Linter.Config} */ 2 | const config = { 3 | extends: [ 4 | "eslint:recommended", 5 | "plugin:@typescript-eslint/recommended-type-checked", 6 | "plugin:@typescript-eslint/stylistic-type-checked", 7 | ], 8 | env: { 9 | es2022: true, 10 | node: true, 11 | }, 12 | parser: "@typescript-eslint/parser", 13 | parserOptions: { project: true }, 14 | plugins: ["@typescript-eslint", "import"], 15 | rules: { 16 | "@typescript-eslint/consistent-type-imports": [ 17 | "warn", 18 | { prefer: "type-imports", fixStyle: "separate-type-imports" }, 19 | ], 20 | "@typescript-eslint/naming-convention": [ 21 | "warn", 22 | { 23 | selector: ["variable", "property"], 24 | format: ["camelCase"], 25 | leadingUnderscore: "allow", 26 | trailingUnderscore: "allow", 27 | }, 28 | { 29 | selector: ["typeLike"], 30 | format: ["PascalCase"], 31 | }, 32 | { 33 | selector: ["variable", "function"], 34 | format: null, 35 | modifiers: ["global"], 36 | }, 37 | ], 38 | "@typescript-eslint/no-misused-promises": [ 39 | 2, 40 | { checksVoidReturn: { attributes: false } }, 41 | ], 42 | "@typescript-eslint/no-unused-vars": [ 43 | "error", 44 | { argsIgnorePattern: "^_", varsIgnorePattern: "^_" }, 45 | ], 46 | "import/consistent-type-specifier-style": ["error", "prefer-top-level"], 47 | }, 48 | ignorePatterns: [".eslintrc.js", "dist"], 49 | reportUnusedDisableDirectives: true, 50 | }; 51 | 52 | module.exports = config; 53 | -------------------------------------------------------------------------------- /.github/workflows/ci_test.yml: -------------------------------------------------------------------------------- 1 | # todo: add tests for Deno and Cloudflare Workers 2 | 3 | name: CI Tests 4 | 5 | on: 6 | push: 7 | branches: [main] 8 | pull_request: 9 | 10 | env: 11 | TURBOPUFFER_API_KEY: ${{ secrets.TURBOPUFFER_API_KEY }} 12 | 13 | jobs: 14 | test-node: 15 | name: Node 16 | runs-on: ubuntu-latest 17 | steps: 18 | - name: Checkout repository 19 | uses: actions/checkout@v3 20 | 21 | - name: Setup Node 22 | uses: actions/setup-node@v3 23 | with: 24 | node-version: "20" 25 | 26 | - name: Install dependencies 27 | run: npm ci 28 | 29 | - name: Run Node tests 30 | run: npm run test:node 31 | 32 | test-bun: 33 | name: Bun 34 | runs-on: ubuntu-latest 35 | needs: test-node 36 | steps: 37 | - name: Checkout repository 38 | uses: actions/checkout@v3 39 | 40 | - name: Install Bun 41 | run: | 42 | curl -fsSL https://bun.sh/install | bash 43 | echo "${HOME}/.bun/bin" >> $GITHUB_PATH 44 | 45 | - name: Install dependencies 46 | run: bun install 47 | 48 | - name: Run Bun tests 49 | run: bun run test:node 50 | 51 | test-browser: 52 | name: Browser (jsdom) 53 | runs-on: ubuntu-latest 54 | needs: test-bun 55 | steps: 56 | - name: Checkout repository 57 | uses: actions/checkout@v3 58 | 59 | - name: Setup Node 60 | uses: actions/setup-node@v3 61 | with: 62 | node-version: "20" 63 | 64 | - name: Install dependencies 65 | run: npm ci 66 | 67 | - name: Run Browser (jsdom) tests 68 | run: npm run test:browser 69 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | node_modules/ 2 | .DS_Store 3 | .env 4 | dist/ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 turbopuffer 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | The **official TypeScript SDK** for turbopuffer. 2 | 3 | To install, 4 | 5 | ```bash 6 | npm i @turbopuffer/turbopuffer 7 | ``` 8 | 9 | Usage: 10 | 11 | ```ts 12 | // Make a new client 13 | // Connections are pooled for the lifetime of the client 14 | // We recommend creating a single instance and reusing it for all calls 15 | const tpuf = new Turbopuffer({ 16 | apiKey: process.env.TURBOPUFFER_API_KEY as string, 17 | // see https://turbopuffer.com/docs/regions for available regions 18 | baseUrl: "https://gcp-us-east4.turbopuffer.com", 19 | }); 20 | 21 | // Instantiate an object to work with a namespace 22 | const ns = tpuf.namespace("readme"); 23 | 24 | await ns.write({ 25 | upsert_rows: [ 26 | { 27 | id: 1, 28 | vector: [1, 2], 29 | foo: "bar", 30 | numbers: [1, 2, 3], 31 | }, 32 | { 33 | id: 2, 34 | vector: [3, 4], 35 | foo: "baz", 36 | numbers: [2, 3, 4], 37 | }, 38 | ], 39 | distance_metric: "cosine_distance", 40 | }); 41 | 42 | const results = await ns.query({ 43 | rank_by: ["vector", "ANN", [1, 1]], 44 | filters: ["numbers", "In", [2, 4]], 45 | top_k: 10, 46 | }); 47 | 48 | // results: 49 | // { 50 | // rows: [ 51 | // { id: 2, $dist: 0.010050535 }, 52 | // { id: 1, $dist: 0.051316738 }, 53 | // ], 54 | // billing: {...}, 55 | // performance: {...} 56 | // } 57 | ``` 58 | 59 | To run the tests, 60 | 61 | ```bash 62 | npm run test 63 | ``` 64 | 65 | To publish a new version, 66 | 67 | 1. Bump version in `package.json` 68 | 2. `npm install` to update `package-lock.json` 69 | 3. `npm publish --access public` 70 | 4. `git tag vX.Y.Z` 71 | 5. `git push origin vX.Y.Z` 72 | -------------------------------------------------------------------------------- /jest_jsdom.config.js: -------------------------------------------------------------------------------- 1 | /** @type {import('ts-jest').JestConfigWithTsJest} */ 2 | module.exports = { 3 | preset: "ts-jest", 4 | testEnvironment: "jsdom", 5 | globals: { 6 | fetch: global.fetch, 7 | }, 8 | }; 9 | -------------------------------------------------------------------------------- /jest_node.config.js: -------------------------------------------------------------------------------- 1 | /** @type {import('ts-jest').JestConfigWithTsJest} */ 2 | module.exports = { 3 | preset: "ts-jest", 4 | testEnvironment: "node", 5 | }; 6 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "@turbopuffer/turbopuffer", 3 | "version": "0.9.1", 4 | "description": "Official Typescript API client library for turbopuffer.com", 5 | "scripts": { 6 | "build": "rm -rf dist && tsc", 7 | "prepublishOnly": "npm run build", 8 | "postinstall:workspaces": "npm run build", 9 | "test:node": "jest --config jest_node.config.js", 10 | "test:browser": "jest --config jest_jsdom.config.js", 11 | "test": "npm run test:node && npm run test:browser", 12 | "format": "prettier --check . --ignore-path ./.gitignore", 13 | "format:fix": "prettier --check . --ignore-path ./.gitignore --write", 14 | "lint": "eslint .", 15 | "lint:fix": "eslint . --fix" 16 | }, 17 | "homepage": "https://turbopuffer.com", 18 | "repository": { 19 | "type": "git", 20 | "url": "git+https://github.com/turbopuffer/turbopuffer-typescript.git" 21 | }, 22 | "browser": { 23 | "./dist/httpClient/node.js": false 24 | }, 25 | "author": "Morgan Gallant ", 26 | "license": "MIT", 27 | "keywords": [ 28 | "turbopuffer", 29 | "embeddings", 30 | "vector database", 31 | "semantic search" 32 | ], 33 | "files": [ 34 | "dist", 35 | "README.md", 36 | "LICENSE" 37 | ], 38 | "main": "./dist/index.js", 39 | "types": "./dist/index.d.ts", 40 | "devDependencies": { 41 | "@types/jest": "^29.5.12", 42 | "@types/pako": "^2.0.3", 43 | "@typescript-eslint/eslint-plugin": "^7.0.2", 44 | "@typescript-eslint/parser": "^7.0.2", 45 | "eslint": "^8.57.0", 46 | "eslint-plugin-import": "^2.29.1", 47 | "jest": "^29.7.0", 48 | "jest-environment-jsdom": "^29.7.0", 49 | "prettier": "^3.3.3", 50 | "ts-jest": "^29.1.2", 51 | "typescript": "^5.4.3" 52 | }, 53 | "dependencies": { 54 | "pako": "^2.1.0", 55 | "undici": "^6.19.8" 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /src/createHTTPClient.ts: -------------------------------------------------------------------------------- 1 | import { isRuntimeFullyNodeCompatible } from "./helpers"; 2 | 3 | /** 4 | * This a helper function that returns a class for making fetch requests 5 | * against the API. 6 | * 7 | * @param baseUrl The base URL of the API endpoint. 8 | * @param apiKey The API key to use for authentication. 9 | * 10 | * @returns An HTTPClient to make requests against the API. 11 | */ 12 | export const createHTTPClient = ( 13 | baseUrl: string, 14 | apiKey: string, 15 | connectTimeout: number, 16 | idleTimeout: number, 17 | warmConnections: number, 18 | compression: boolean, 19 | ) => { 20 | if (isRuntimeFullyNodeCompatible) { 21 | // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment,@typescript-eslint/no-var-requires, @typescript-eslint/no-unsafe-member-access 22 | const NodeHTTPClient = require("./httpClient/node").default; 23 | // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-call 24 | return new NodeHTTPClient( 25 | baseUrl, 26 | apiKey, 27 | connectTimeout, 28 | idleTimeout, 29 | warmConnections, 30 | compression, 31 | ); 32 | } else { 33 | // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-var-requires, @typescript-eslint/no-unsafe-member-access 34 | const DefaultHTTPClient = require("./httpClient/default").default; 35 | // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-call 36 | return new DefaultHTTPClient(baseUrl, apiKey, warmConnections, compression); 37 | } 38 | }; 39 | -------------------------------------------------------------------------------- /src/helpers.ts: -------------------------------------------------------------------------------- 1 | import type { RequestTiming, WriteParams } from "./types"; 2 | 3 | type Runtime = 4 | | "bun" 5 | | "deno" 6 | | "cloudflare-workers" 7 | | "browser" 8 | | "node" 9 | | undefined; 10 | 11 | function detectRuntime(): Runtime { 12 | // @ts-expect-error can be ignored 13 | if (typeof globalThis.Bun !== "undefined") return "bun"; 14 | 15 | // @ts-expect-error can be ignored 16 | if (typeof globalThis.Deno !== "undefined") return "deno"; 17 | 18 | const userAgent = globalThis.navigator?.userAgent; 19 | 20 | // Try navigator.userAgent: 21 | // https://developers.cloudflare.com/workers/runtime-apis/web-standards/#navigatoruseragent. 22 | // Fallback: look for presence of non-standard globals specific to the cloudflare runtime: 23 | // https://community.cloudflare.com/t/how-to-detect-the-cloudflare-worker-runtime/293715/2. 24 | if ( 25 | userAgent 26 | ? userAgent === "Cloudflare-Workers" 27 | : // @ts-expect-error can be ignored 28 | typeof WebSocketPair !== "undefined" 29 | ) 30 | return "cloudflare-workers"; 31 | 32 | if (typeof window !== "undefined") return "browser"; 33 | 34 | if ( 35 | userAgent 36 | ? userAgent.startsWith("Node.js") 37 | : process.release?.name === "node" 38 | ) 39 | return "node"; 40 | } 41 | 42 | const detectedRuntime = detectRuntime(); 43 | export const isRuntimeFullyNodeCompatible = 44 | detectedRuntime === "node" || detectedRuntime === "deno"; 45 | 46 | /** An error class for errors returned by the turbopuffer API. */ 47 | export class TurbopufferError extends Error { 48 | status?: number; 49 | constructor( 50 | public error: string, 51 | { status, cause }: { status?: number; cause?: Error }, 52 | ) { 53 | super(error, { cause: cause }); 54 | this.status = status; 55 | } 56 | } 57 | 58 | export function buildBaseUrl(baseUrl: string) { 59 | const url = baseUrl.trim(); 60 | const hasProtocol = url.includes("://"); 61 | return hasProtocol ? url : `https://${url}`; 62 | } 63 | 64 | export function buildUrl( 65 | baseUrl: string, 66 | path: string, 67 | query?: Record, 68 | ) { 69 | // https://developer.mozilla.org/en-US/docs/Web/API/URL_API/Resolving_relative_references 70 | // if baseUrl doesn't end in /, add one to make it behave 71 | // like a directory so the next path part is appended. 72 | // if there are multiple / appended, ensure all but one get removed. 73 | const updatedBaseUrl = baseUrl.replace(/\/*$/, "/"); 74 | 75 | // strip leading slashes from `path` so it's appended rather 76 | // than treated as absolute 77 | const relativePath = path.replace(/^\/+/, ""); 78 | 79 | const url = new URL(relativePath, updatedBaseUrl); 80 | 81 | if (query) { 82 | for (const [key, value] of Object.entries(query)) { 83 | if (value) url.searchParams.set(key, value); 84 | } 85 | } 86 | 87 | return url; 88 | } 89 | 90 | /** A helper function to determine if a status code should be retried. */ 91 | export function statusCodeShouldRetry(statusCode?: number): boolean { 92 | return ( 93 | !statusCode || statusCode === 408 || statusCode === 409 || statusCode === 429 || statusCode >= 500 94 | ); 95 | } 96 | 97 | /** A helper function to delay for a given number of milliseconds. */ 98 | export function delay(ms: number) { 99 | return new Promise((resolve) => setTimeout(resolve, ms)); 100 | } 101 | 102 | export function make_request_timing({ 103 | request_start, 104 | response_start, 105 | body_read_end, 106 | decompress_end, 107 | deserialize_end, 108 | requestCompressionDuration, 109 | }: { 110 | request_start: number; 111 | response_start: number; 112 | body_read_end?: number; 113 | decompress_end?: number; 114 | deserialize_end?: number; 115 | requestCompressionDuration?: number; 116 | }): RequestTiming { 117 | const deserialize_start = decompress_end ?? body_read_end; 118 | return { 119 | response_time: response_start - request_start, 120 | // `!= null` checks for both null and undefined 121 | body_read_time: 122 | body_read_end != null ? body_read_end - response_start : null, 123 | compress_time: requestCompressionDuration ?? null, 124 | decompress_time: 125 | decompress_end != null && body_read_end != null 126 | ? decompress_end - body_read_end 127 | : null, 128 | deserialize_time: 129 | deserialize_end != null && deserialize_start != null 130 | ? deserialize_end - deserialize_start 131 | : null, 132 | }; 133 | } 134 | 135 | export function shouldCompressWrite({ 136 | upsert_columns, 137 | upsert_rows, 138 | patch_columns, 139 | patch_rows, 140 | deletes, 141 | }: WriteParams): boolean { 142 | return ( 143 | (upsert_columns?.id.length ?? 0) > 10 || 144 | (upsert_rows?.length ?? 0) > 10 || 145 | (upsert_rows?.some((row) => (row.vector?.length ?? 0) > 10) ?? false) || 146 | (patch_columns?.id.length ?? 0) > 10 || 147 | (patch_rows?.length ?? 0) > 10 || 148 | (deletes?.length ?? 0) > 500 149 | ); 150 | } 151 | -------------------------------------------------------------------------------- /src/httpClient/default.ts: -------------------------------------------------------------------------------- 1 | import { gzip } from "pako"; 2 | import { version } from "../../package.json"; 3 | import type { RequestParams, RequestResponse, HTTPClient } from "../types"; 4 | import { 5 | TurbopufferError, 6 | statusCodeShouldRetry, 7 | delay, 8 | make_request_timing, 9 | buildBaseUrl, 10 | buildUrl, 11 | } from "../helpers"; 12 | 13 | function convertHeadersType(headers: Headers): Record { 14 | const normalizedHeaders: Record = {}; 15 | for (const [key, value] of headers.entries()) { 16 | if (value === undefined) { 17 | continue; 18 | } else if (Array.isArray(value)) { 19 | normalizedHeaders[key] = value[0] as string; 20 | } else { 21 | normalizedHeaders[key] = value; 22 | } 23 | } 24 | return normalizedHeaders; 25 | } 26 | 27 | export default class DefaultHTTPClient implements HTTPClient { 28 | private baseUrl: string; 29 | private origin: URL; 30 | private apiKey: string; 31 | readonly userAgent = `tpuf-typescript/${version}/fetch`; 32 | private compression: boolean; 33 | 34 | constructor( 35 | baseUrl: string, 36 | apiKey: string, 37 | warmConnections: number, 38 | compression: boolean, 39 | ) { 40 | this.baseUrl = buildBaseUrl(baseUrl); 41 | this.origin = new URL(baseUrl); 42 | this.origin.pathname = ""; 43 | this.apiKey = apiKey; 44 | this.compression = compression; 45 | 46 | for (let i = 0; i < warmConnections; i++) { 47 | // send a small request to put some connections in the pool 48 | void fetch(this.baseUrl, { 49 | method: "HEAD", 50 | headers: { "User-Agent": this.userAgent }, 51 | }); 52 | } 53 | } 54 | 55 | async doRequest({ 56 | method, 57 | path, 58 | query, 59 | body, 60 | compress, 61 | retryable, 62 | }: RequestParams): RequestResponse { 63 | const url = buildUrl(this.baseUrl, path, query); 64 | 65 | const headers: Record = { 66 | Authorization: `Bearer ${this.apiKey}`, 67 | "User-Agent": this.userAgent, 68 | }; 69 | 70 | if (this.compression) { 71 | headers["Accept-Encoding"] = "gzip"; 72 | } 73 | 74 | if (body) { 75 | headers["Content-Type"] = "application/json"; 76 | } 77 | 78 | let requestCompressionDuration; 79 | let requestBody: Uint8Array | string | null = null; 80 | if (body && compress && this.compression) { 81 | headers["Content-Encoding"] = "gzip"; 82 | const beforeRequestCompression = performance.now(); 83 | requestBody = gzip(JSON.stringify(body)); 84 | requestCompressionDuration = performance.now() - beforeRequestCompression; 85 | } else if (body) { 86 | requestBody = JSON.stringify(body); 87 | } 88 | 89 | const maxAttempts = retryable ? 3 : 1; 90 | let response!: Response; 91 | let error: TurbopufferError | null = null; 92 | let request_start!: number; 93 | let response_start!: number; 94 | 95 | for (let attempt = 0; attempt < maxAttempts; attempt++) { 96 | error = null; 97 | request_start = performance.now(); 98 | try { 99 | response = await fetch(url, { 100 | method, 101 | headers, 102 | body: requestBody, 103 | }); 104 | } catch (e: unknown) { 105 | if (e instanceof Error) { 106 | error = new TurbopufferError(`fetch failed: ${e.message}`, { 107 | cause: e, 108 | }); 109 | } else { 110 | // not an Error? shouldn't happen but good to be thorough 111 | throw e; 112 | } 113 | } 114 | response_start = performance.now(); 115 | 116 | if (!error && response.status >= 400) { 117 | let message: string | undefined = undefined; 118 | const body_text = await response.text(); 119 | if (response.headers.get("content-type") === "application/json") { 120 | try { 121 | const body = JSON.parse(body_text); 122 | if (body && body.status === "error") { 123 | message = body.error; 124 | } else { 125 | message = body_text; 126 | } 127 | } catch (_: unknown) { 128 | /* empty */ 129 | } 130 | } else { 131 | message = body_text; 132 | } 133 | error = new TurbopufferError( 134 | message ?? `http error ${response.status}`, 135 | { 136 | status: response.status, 137 | }, 138 | ); 139 | } 140 | if ( 141 | error && 142 | statusCodeShouldRetry(error.status) && 143 | attempt + 1 !== maxAttempts 144 | ) { 145 | await delay(150 * (attempt + 1)); // 150ms, 300ms, 450ms 146 | continue; 147 | } 148 | break; 149 | } 150 | if (error) { 151 | throw error; 152 | } 153 | 154 | if (method === "HEAD" || !response.body) { 155 | return { 156 | headers: convertHeadersType(response.headers), 157 | request_timing: make_request_timing({ request_start, response_start }), 158 | }; 159 | } 160 | 161 | const body_text = await response.text(); 162 | const body_read_end = performance.now(); 163 | 164 | const json = JSON.parse(body_text); 165 | const deserialize_end = performance.now(); 166 | 167 | if (json.status && json.status === "error") { 168 | throw new TurbopufferError(json.error || (json as string), { 169 | status: response.status, 170 | }); 171 | } 172 | 173 | return { 174 | body: json as T, 175 | headers: convertHeadersType(response.headers), 176 | request_timing: make_request_timing({ 177 | request_start, 178 | response_start, 179 | body_read_end, 180 | deserialize_end, 181 | requestCompressionDuration, 182 | }), 183 | }; 184 | } 185 | } 186 | -------------------------------------------------------------------------------- /src/httpClient/node.ts: -------------------------------------------------------------------------------- 1 | import { fetch, Agent } from "undici"; 2 | import { gzip, gunzip } from "node:zlib"; 3 | import { promisify } from "node:util"; 4 | import type { Dispatcher } from "undici"; 5 | 6 | import { version } from "../../package.json"; 7 | import type { 8 | RequestParams, 9 | RequestResponse, 10 | HTTPClient, 11 | TpufResponseWithMetadata, 12 | } from "../types"; 13 | import { 14 | TurbopufferError, 15 | statusCodeShouldRetry, 16 | delay, 17 | make_request_timing, 18 | buildBaseUrl, 19 | buildUrl, 20 | } from "../helpers"; 21 | 22 | const gzipAsync = promisify(gzip); 23 | const gunzipAsync = promisify(gunzip); 24 | 25 | function convertHeadersType( 26 | headers: Record 27 | ): Record { 28 | for (const key in headers) { 29 | const v = headers[key]; 30 | if (v === undefined) { 31 | delete headers[key]; 32 | } else if (Array.isArray(v)) { 33 | headers[key] = v[0]; 34 | } 35 | } 36 | return headers as Record; 37 | } 38 | 39 | async function consumeResponseText( 40 | response: Dispatcher.ResponseData 41 | ): Promise { 42 | if (response.headers["content-encoding"] === "gzip") { 43 | const body_buffer = await response.body.arrayBuffer(); 44 | const body_read_end = performance.now(); 45 | 46 | const gunzip_buffer = await gunzipAsync(body_buffer); 47 | const body_text = gunzip_buffer.toString(); // is there a better way? 48 | const decompress_end = performance.now(); 49 | return { body_text, body_read_end, decompress_end }; 50 | } else { 51 | const body_text = await response.body.text(); 52 | const body_read_end = performance.now(); 53 | return { body_text, body_read_end, decompress_end: body_read_end }; 54 | } 55 | } 56 | 57 | export default class NodeHTTPClient implements HTTPClient { 58 | private agent: Agent; 59 | private baseUrl: string; 60 | private origin: URL; 61 | private apiKey: string; 62 | readonly userAgent = `tpuf-typescript/${version}/node`; 63 | private compression: boolean; 64 | 65 | constructor( 66 | baseUrl: string, 67 | apiKey: string, 68 | connectTimeout: number, 69 | idleTimeout: number, 70 | warmConnections: number, 71 | compression: boolean 72 | ) { 73 | this.baseUrl = buildBaseUrl(baseUrl); 74 | this.origin = new URL(this.baseUrl); 75 | this.origin.pathname = ""; 76 | this.apiKey = apiKey; 77 | this.compression = compression; 78 | 79 | this.agent = new Agent({ 80 | keepAliveTimeout: idleTimeout, // how long a socket can be idle for before it is closed 81 | keepAliveMaxTimeout: 24 * 60 * 60 * 1000, // maximum configurable timeout with server hint 82 | connect: { 83 | timeout: connectTimeout, 84 | }, 85 | }); 86 | 87 | for (let i = 0; i < warmConnections; i++) { 88 | // send a small request to put some connections in the pool 89 | void fetch(this.baseUrl, { 90 | method: "HEAD", 91 | headers: { "User-Agent": this.userAgent }, 92 | dispatcher: this.agent, 93 | }); 94 | } 95 | } 96 | 97 | async doRequest({ 98 | method, 99 | path, 100 | query, 101 | body, 102 | compress, 103 | retryable, 104 | }: RequestParams): RequestResponse { 105 | const url = buildUrl(this.baseUrl, path, query); 106 | 107 | const headers: Record = { 108 | Authorization: `Bearer ${this.apiKey}`, 109 | "User-Agent": this.userAgent, 110 | }; 111 | 112 | if (this.compression) { 113 | headers["Accept-Encoding"] = "gzip"; 114 | } 115 | 116 | if (body) { 117 | headers["Content-Type"] = "application/json"; 118 | } 119 | 120 | let requestCompressionDuration; 121 | let requestBody: Uint8Array | string | null = null; 122 | if (body && compress && this.compression) { 123 | headers["Content-Encoding"] = "gzip"; 124 | const beforeRequestCompression = performance.now(); 125 | requestBody = await gzipAsync(JSON.stringify(body)); 126 | requestCompressionDuration = performance.now() - beforeRequestCompression; 127 | } else if (body) { 128 | requestBody = JSON.stringify(body); 129 | } 130 | 131 | const maxAttempts = retryable ? 3 : 1; 132 | let response!: Dispatcher.ResponseData; 133 | let error: TurbopufferError | null = null; 134 | let request_start!: number; 135 | let response_start!: number; 136 | 137 | for (let attempt = 0; attempt < maxAttempts; attempt++) { 138 | error = null; 139 | request_start = performance.now(); 140 | try { 141 | response = await this.agent.request({ 142 | origin: url.origin, 143 | path: url.pathname + url.search, 144 | method: method as Dispatcher.HttpMethod, 145 | headers, 146 | body: requestBody, 147 | }); 148 | } catch (e: unknown) { 149 | if (e instanceof Error) { 150 | if (e.cause instanceof Error) { 151 | // wrap generic undici "fetch failed" error with the underlying cause 152 | error = new TurbopufferError(`fetch failed: ${e.cause.message}`, { 153 | cause: e, 154 | }); 155 | } else { 156 | // wrap other errors directly 157 | error = new TurbopufferError(`fetch failed: ${e.message}`, { 158 | cause: e, 159 | }); 160 | } 161 | } else { 162 | // not an Error? shouldn't happen but good to be thorough 163 | throw e; 164 | } 165 | } 166 | response_start = performance.now(); 167 | 168 | if (!error && response.statusCode >= 400) { 169 | let message: string | undefined = undefined; 170 | const { body_text } = await consumeResponseText(response); 171 | if (response.headers["content-type"] === "application/json") { 172 | try { 173 | const body = JSON.parse(body_text); 174 | if (body && body.status === "error") { 175 | message = body.error; 176 | } else { 177 | message = body_text; 178 | } 179 | } catch (_: unknown) { 180 | /* empty */ 181 | } 182 | } else { 183 | message = body_text; 184 | } 185 | error = new TurbopufferError( 186 | message ?? `http error ${response.statusCode}`, 187 | { 188 | status: response.statusCode, 189 | } 190 | ); 191 | } 192 | if ( 193 | error && 194 | statusCodeShouldRetry(error.status) && 195 | attempt + 1 !== maxAttempts 196 | ) { 197 | await delay(150 * (attempt + 1)); // 150ms, 300ms, 450ms 198 | continue; 199 | } 200 | break; 201 | } 202 | if (error) { 203 | throw error; 204 | } 205 | 206 | if (method === "HEAD" || !response.body) { 207 | return { 208 | headers: convertHeadersType(response.headers), 209 | request_timing: make_request_timing({ request_start, response_start }), 210 | }; 211 | } 212 | 213 | const { body_text, body_read_end, decompress_end } = 214 | await consumeResponseText(response); 215 | 216 | const json = JSON.parse(body_text); 217 | const deserialize_end = performance.now(); 218 | 219 | if (json.status && json.status === "error") { 220 | throw new TurbopufferError(json.error || (json as string), { 221 | status: response.statusCode, 222 | }); 223 | } 224 | 225 | return { 226 | body: json as T, 227 | headers: convertHeadersType(response.headers), 228 | request_timing: make_request_timing({ 229 | request_start, 230 | response_start, 231 | body_read_end, 232 | decompress_end, 233 | deserialize_end, 234 | requestCompressionDuration, 235 | }), 236 | }; 237 | } 238 | } 239 | -------------------------------------------------------------------------------- /src/index.ts: -------------------------------------------------------------------------------- 1 | export * from "./types"; 2 | 3 | export { TurbopufferError } from "./helpers"; 4 | export * from "./turbopuffer"; 5 | -------------------------------------------------------------------------------- /src/turbopuffer.test.ts: -------------------------------------------------------------------------------- 1 | import { Turbopuffer, RankBy, Schema, TurbopufferError } from "./index"; 2 | import { isRuntimeFullyNodeCompatible, buildBaseUrl, buildUrl } from "./helpers"; 3 | 4 | const tpuf = new Turbopuffer({ 5 | apiKey: process.env.TURBOPUFFER_API_KEY!, 6 | }); 7 | 8 | const testNamespacePrefix = "typescript_sdk_"; 9 | 10 | test("trailing_slashes_in_base_url", async () => { 11 | const tpuf = new Turbopuffer({ 12 | apiKey: process.env.TURBOPUFFER_API_KEY!, 13 | baseUrl: "https://gcp-us-east4.turbopuffer.com//", 14 | }); 15 | 16 | const ns = tpuf.namespace( 17 | testNamespacePrefix + "trailing_slashes_in_base_url", 18 | ); 19 | 20 | await ns.write({ 21 | upsert_rows: [ 22 | { 23 | id: 1, 24 | vector: [0.1, 0.1], 25 | text: "Walruses are large marine mammals with long tusks and whiskers", 26 | }, 27 | { 28 | id: 2, 29 | vector: [0.2, 0.2], 30 | text: "They primarily inhabit the cold Arctic regions", 31 | }, 32 | ], 33 | distance_metric: "cosine_distance", 34 | }); 35 | 36 | const schema = await ns.schema(); 37 | expect(schema).toEqual({ 38 | id: { 39 | type: "uint", 40 | filterable: null, 41 | full_text_search: null, 42 | }, 43 | text: { 44 | type: "string", 45 | filterable: true, 46 | full_text_search: null, 47 | }, 48 | vector: { 49 | type: "[2]f32", 50 | ann: true, 51 | filterable: null, 52 | full_text_search: null, 53 | }, 54 | }); 55 | }); 56 | 57 | test("bm25_with_custom_schema_and_sum_query", async () => { 58 | const ns = tpuf.namespace( 59 | testNamespacePrefix + "bm25_with_custom_schema_and_sum_query", 60 | ); 61 | 62 | try { 63 | await ns.deleteAll(); 64 | } catch (_: unknown) { 65 | /* empty */ 66 | } 67 | 68 | await ns.write({ 69 | upsert_rows: [ 70 | { 71 | id: 1, 72 | vector: [0.1, 0.1], 73 | text: "Walruses are large marine mammals with long tusks and whiskers", 74 | }, 75 | { 76 | id: 2, 77 | vector: [0.2, 0.2], 78 | text: "They primarily inhabit the cold Arctic regions", 79 | }, 80 | { 81 | id: 3, 82 | vector: [0.3, 0.3], 83 | text: "Walruses use their tusks to help haul themselves onto ice", 84 | }, 85 | { 86 | id: 4, 87 | vector: [0.4, 0.4], 88 | text: "Their diet mainly consists of mollusks and other sea creatures", 89 | }, 90 | { 91 | id: 5, 92 | vector: [0.5, 0.5], 93 | text: "Walrus populations are affected by climate change and melting ice", 94 | }, 95 | ], 96 | distance_metric: "cosine_distance", 97 | schema: { 98 | text: { 99 | type: "string", 100 | full_text_search: { 101 | language: "english", 102 | stemming: true, 103 | case_sensitive: false, 104 | remove_stopwords: true, 105 | }, 106 | }, 107 | }, 108 | }); 109 | 110 | const results = await ns.query({ 111 | rank_by: [ 112 | "Sum", 113 | [ 114 | ["text", "BM25", "large tusk"], 115 | ["text", "BM25", "mollusk diet"], 116 | ], 117 | ], 118 | top_k: 10, 119 | }); 120 | 121 | expect(results.rows.length).toEqual(3); 122 | expect(results.rows[0].id).toEqual(4); 123 | expect(results.rows[1].id).toEqual(1); 124 | expect(results.rows[2].id).toEqual(3); 125 | }); 126 | 127 | test("bm25_with_tokenizer_pre_tokenized_array", async () => { 128 | const ns = tpuf.namespace( 129 | testNamespacePrefix + "bm25_with_tokenizer_pre_tokenized_array", 130 | ); 131 | try { 132 | await ns.deleteAll(); 133 | } catch (_: unknown) { 134 | /* empty */ 135 | } 136 | 137 | // let's test with a columnar write 138 | await ns.write({ 139 | upsert_columns: { 140 | id: [1, 2], 141 | vector: [ 142 | [0.1, 0.1], 143 | [0.2, 0.2], 144 | ], 145 | content: [ 146 | ["jumped", "over", "the", "lazy", "dog"], 147 | ["the", "lazy", "dog", "is", "brown"], 148 | ], 149 | }, 150 | schema: { 151 | content: { 152 | type: "[]string", 153 | full_text_search: { 154 | tokenizer: "pre_tokenized_array", 155 | }, 156 | }, 157 | }, 158 | distance_metric: "cosine_distance", 159 | }); 160 | 161 | let results = await ns.query({ 162 | rank_by: ["content", "BM25", ["jumped", "over"]], 163 | top_k: 10, 164 | }); 165 | expect(results.rows.length).toEqual(1); 166 | expect(results.rows[0].id).toEqual(1); 167 | 168 | results = await ns.query({ 169 | rank_by: ["content", "BM25", ["dog"]], 170 | top_k: 10, 171 | }); 172 | expect(results.rows.length).toEqual(2); 173 | 174 | await expect( 175 | ns.query({ 176 | rank_by: ["content", "BM25", "jumped"], 177 | top_k: 10, 178 | }), 179 | ).rejects.toThrow( 180 | "invalid input 'jumped' for rank_by field \"content\", expecting []string", 181 | ); 182 | }); 183 | 184 | test("contains_all_tokens", async () => { 185 | const ns = tpuf.namespace(testNamespacePrefix + "contains_all_tokens"); 186 | try { 187 | await ns.deleteAll(); 188 | } catch (_: unknown) { 189 | /* empty */ 190 | } 191 | 192 | // let's test with a columnar write 193 | await ns.write({ 194 | upsert_columns: { 195 | id: [1], 196 | vector: [[0.1, 0.1]], 197 | text: ["Walruses are large marine mammals with long tusks and whiskers"], 198 | }, 199 | schema: { 200 | text: { 201 | type: "string", 202 | full_text_search: { 203 | stemming: true, 204 | }, 205 | }, 206 | }, 207 | distance_metric: "cosine_distance", 208 | }); 209 | 210 | const results = await ns.query({ 211 | rank_by: ["text", "BM25", "walrus whisker"], 212 | filters: ["text", "ContainsAllTokens", "marine mammals"], 213 | top_k: 10, 214 | }); 215 | expect(results.rows.length).toEqual(1); 216 | 217 | const missing = await ns.query({ 218 | rank_by: ["text", "BM25", "walrus whisker"], 219 | filters: ["text", "ContainsAllTokens", "marine mammals short"], 220 | top_k: 10, 221 | }); 222 | expect(missing.rows.length).toEqual(0); 223 | }); 224 | 225 | test("order_by_attribute", async () => { 226 | const ns = tpuf.namespace(testNamespacePrefix + "order_by_attribute"); 227 | 228 | try { 229 | await ns.deleteAll(); 230 | } catch (_: unknown) { 231 | /* empty */ 232 | } 233 | 234 | await ns.write({ 235 | upsert_rows: [ 236 | { 237 | id: 1, 238 | vector: [0.1, 0.1], 239 | a: "5", 240 | }, 241 | { 242 | id: 2, 243 | vector: [0.2, 0.2], 244 | a: "4", 245 | }, 246 | { 247 | id: 3, 248 | vector: [0.3, 0.3], 249 | a: "3", 250 | }, 251 | { 252 | id: 4, 253 | vector: [0.4, 0.4], 254 | a: "2", 255 | }, 256 | { 257 | id: 5, 258 | vector: [0.5, 0.5], 259 | a: "1", 260 | }, 261 | ], 262 | distance_metric: "euclidean_squared", 263 | }); 264 | 265 | const results_asc = await ns.query({ 266 | rank_by: ["a", "asc"], 267 | top_k: 10, 268 | }); 269 | expect(results_asc.rows.length).toEqual(5); 270 | expect(results_asc.rows[0].id).toEqual(5); 271 | expect(results_asc.rows[1].id).toEqual(4); 272 | expect(results_asc.rows[2].id).toEqual(3); 273 | expect(results_asc.rows[3].id).toEqual(2); 274 | expect(results_asc.rows[4].id).toEqual(1); 275 | 276 | const results_desc = await ns.query({ 277 | rank_by: ["a", "desc"], 278 | top_k: 10, 279 | }); 280 | expect(results_desc.rows.length).toEqual(5); 281 | expect(results_desc.rows[0].id).toEqual(1); 282 | expect(results_desc.rows[1].id).toEqual(2); 283 | expect(results_desc.rows[2].id).toEqual(3); 284 | expect(results_desc.rows[3].id).toEqual(4); 285 | expect(results_desc.rows[4].id).toEqual(5); 286 | }); 287 | 288 | test("bm25_with_default_schema_and_simple_query", async () => { 289 | const ns = tpuf.namespace( 290 | testNamespacePrefix + "bm25_with_default_schema_and_simple_query", 291 | ); 292 | 293 | try { 294 | await ns.deleteAll(); 295 | } catch (_: unknown) { 296 | /* empty */ 297 | } 298 | 299 | await ns.write({ 300 | upsert_rows: [ 301 | { 302 | id: 1, 303 | vector: [0.1, 0.1], 304 | text: "Walruses can produce a variety of funny sounds, including whistles, grunts, and bell-like noises.", 305 | }, 306 | { 307 | id: 2, 308 | vector: [0.2, 0.2], 309 | text: "They sometimes use their tusks as a tool to break through ice or to scratch their bodies.", 310 | }, 311 | ], 312 | distance_metric: "cosine_distance", 313 | schema: { 314 | text: { 315 | type: "string", 316 | full_text_search: true, 317 | }, 318 | }, 319 | }); 320 | 321 | const results = await ns.query({ 322 | rank_by: ["text", "BM25", "scratch"], 323 | top_k: 10, 324 | }); 325 | 326 | expect(results.rows.length).toEqual(1); 327 | expect(results.rows[0].id).toEqual(2); 328 | }); 329 | 330 | test("namespaces", async () => { 331 | const namespaces0 = await tpuf.namespaces({ page_size: 5 }); 332 | const cursor0 = namespaces0.next_cursor; 333 | 334 | const namespaces1 = await tpuf.namespaces({ 335 | cursor: cursor0, 336 | page_size: 5, 337 | }); 338 | const cursor1 = namespaces1.next_cursor; 339 | 340 | expect(namespaces0.namespaces.length).toEqual(5); 341 | expect(namespaces0.namespaces.length).toEqual(5); 342 | expect(cursor0).not.toEqual(cursor1); 343 | }); 344 | 345 | test("hint_cache_warm", async () => { 346 | const nsId = (await tpuf.namespaces({ page_size: 1 })).namespaces[0].id; 347 | const ns = await tpuf.namespace(nsId); 348 | 349 | const result = await ns.hintCacheWarm(); 350 | 351 | expect(typeof result.message).toBe("string"); 352 | expect(["ACCEPTED", "OK"]).toContain(result.status); 353 | }); 354 | 355 | test("schema", async () => { 356 | const ns = tpuf.namespace(testNamespacePrefix + "schema"); 357 | 358 | try { 359 | await ns.deleteAll(); 360 | } catch (_: unknown) { 361 | /* empty */ 362 | } 363 | 364 | await ns.write({ 365 | upsert_rows: [ 366 | { 367 | id: 1, 368 | vector: [0.1, 0.1], 369 | title: "one", 370 | private: true, 371 | tags: ["a", "b"], 372 | }, 373 | { 374 | id: 2, 375 | vector: [0.2, 0.2], 376 | title: null, 377 | private: null, 378 | tags: ["b", "d"], 379 | }, 380 | { 381 | id: 3, 382 | vector: [0.3, 0.3], 383 | title: "three", 384 | private: false, 385 | tags: [], 386 | }, 387 | { 388 | id: 4, 389 | vector: [0.4, 0.4], 390 | title: "four", 391 | private: true, 392 | tags: ["c"], 393 | }, 394 | ], 395 | distance_metric: "cosine_distance", 396 | schema: { 397 | title: { 398 | type: "string", 399 | full_text_search: { 400 | stemming: true, 401 | remove_stopwords: true, 402 | case_sensitive: false, 403 | }, 404 | }, 405 | tags: { 406 | type: "[]string", 407 | full_text_search: { 408 | stemming: false, 409 | remove_stopwords: false, 410 | case_sensitive: true, 411 | }, 412 | }, 413 | vector: { 414 | type: "[2]f16", 415 | ann: true, 416 | }, 417 | }, 418 | }); 419 | 420 | const schema = await ns.schema(); 421 | expect(schema).toEqual({ 422 | id: { 423 | type: "uint", 424 | filterable: null, 425 | full_text_search: null, 426 | }, 427 | title: { 428 | type: "string", 429 | filterable: false, 430 | full_text_search: { 431 | k1: 1.2, 432 | b: 0.75, 433 | language: "english", 434 | stemming: true, 435 | remove_stopwords: true, 436 | case_sensitive: false, 437 | tokenizer: "word_v1", 438 | }, 439 | }, 440 | tags: { 441 | type: "[]string", 442 | filterable: false, 443 | full_text_search: { 444 | k1: 1.2, 445 | b: 0.75, 446 | language: "english", 447 | stemming: false, 448 | remove_stopwords: false, 449 | case_sensitive: true, 450 | tokenizer: "word_v1", 451 | }, 452 | }, 453 | private: { 454 | type: "bool", 455 | filterable: true, 456 | full_text_search: null, 457 | }, 458 | vector: { 459 | type: "[2]f16", 460 | ann: true, 461 | filterable: null, 462 | full_text_search: null, 463 | }, 464 | }); 465 | }); 466 | 467 | test("update_schema", async () => { 468 | const ns = tpuf.namespace(testNamespacePrefix + "update_schema"); 469 | 470 | try { 471 | await ns.deleteAll(); 472 | } catch (_: unknown) { 473 | /* empty */ 474 | } 475 | 476 | await ns.write({ 477 | upsert_rows: [ 478 | { 479 | id: 1, 480 | vector: [0.1, 0.1], 481 | private: true, 482 | tags: ["a", "b"], 483 | }, 484 | { 485 | id: 2, 486 | vector: [0.2, 0.2], 487 | private: null, 488 | tags: ["b", "d"], 489 | }, 490 | ], 491 | distance_metric: "cosine_distance", 492 | schema: { 493 | tags: { 494 | type: "[]string", 495 | full_text_search: { 496 | stemming: false, 497 | remove_stopwords: false, 498 | case_sensitive: true, 499 | }, 500 | }, 501 | }, 502 | }); 503 | 504 | const schema = await ns.schema(); 505 | expect(schema).toEqual({ 506 | id: { 507 | type: "uint", 508 | filterable: null, 509 | full_text_search: null, 510 | }, 511 | tags: { 512 | type: "[]string", 513 | filterable: false, 514 | full_text_search: { 515 | k1: 1.2, 516 | b: 0.75, 517 | language: "english", 518 | stemming: false, 519 | remove_stopwords: false, 520 | case_sensitive: true, 521 | tokenizer: "word_v1", 522 | }, 523 | }, 524 | private: { 525 | type: "bool", 526 | filterable: true, 527 | full_text_search: null, 528 | }, 529 | vector: { 530 | type: "[2]f32", 531 | ann: true, 532 | filterable: null, 533 | full_text_search: null, 534 | }, 535 | }); 536 | 537 | // Write an update to the schema making 'tags' 538 | // filterable and 'private' not filterable 539 | const updateSchema = await ns.updateSchema({ 540 | tags: { 541 | type: "[]string", 542 | filterable: true, 543 | full_text_search: { 544 | k1: 1.2, 545 | b: 0.75, 546 | language: "english", 547 | stemming: false, 548 | remove_stopwords: false, 549 | case_sensitive: true, 550 | tokenizer: "word_v1", 551 | }, 552 | }, 553 | private: { 554 | type: "bool", 555 | filterable: false, 556 | full_text_search: false, 557 | }, 558 | }); 559 | expect(updateSchema).toEqual({ 560 | id: { 561 | type: "uint", 562 | filterable: null, 563 | full_text_search: null, 564 | }, 565 | tags: { 566 | type: "[]string", 567 | filterable: true, 568 | full_text_search: { 569 | k1: 1.2, 570 | b: 0.75, 571 | language: "english", 572 | stemming: false, 573 | remove_stopwords: false, 574 | case_sensitive: true, 575 | tokenizer: "word_v1", 576 | }, 577 | }, 578 | private: { 579 | type: "bool", 580 | filterable: false, 581 | full_text_search: null, 582 | }, 583 | vector: { 584 | type: "[2]f32", 585 | ann: true, 586 | filterable: null, 587 | full_text_search: null, 588 | }, 589 | }); 590 | }); 591 | 592 | test("sanity", async () => { 593 | const nameSpaceName = testNamespacePrefix + "sanity"; 594 | const ns = tpuf.namespace(nameSpaceName); 595 | 596 | try { 597 | await ns.deleteAll(); 598 | } catch (_: unknown) { 599 | /* empty */ 600 | } 601 | 602 | await ns.write({ 603 | upsert_rows: [ 604 | { 605 | id: 1, 606 | vector: [1, 2], 607 | foo: "bar", 608 | numbers: [1, 2, 3], 609 | maybeNull: null, 610 | bool: true, 611 | }, 612 | { 613 | id: 2, 614 | vector: [3, 4], 615 | foo: "baz", 616 | numbers: [2, 3, 4], 617 | maybeNull: null, 618 | bool: true, 619 | }, 620 | { 621 | id: 3, 622 | vector: [3, 4], 623 | foo: "baz", 624 | numbers: [17], 625 | maybeNull: "oh boy!", 626 | bool: true, 627 | }, 628 | ], 629 | distance_metric: "cosine_distance", 630 | }); 631 | 632 | const resultsWithPerformance = await ns.query({ 633 | rank_by: ["vector", "ANN", [1, 1]], 634 | filters: ["numbers", "In", [2, 4]], 635 | top_k: 10, 636 | }); 637 | expect(resultsWithPerformance.rows.length).toEqual(2); 638 | expect(resultsWithPerformance.rows[0].id).toEqual(2); 639 | expect(resultsWithPerformance.rows[1].id).toEqual(1); 640 | 641 | const performance = resultsWithPerformance.performance; 642 | expect(performance.approx_namespace_size).toEqual(3); 643 | expect(performance.exhaustive_search_count).toEqual(3); 644 | expect(performance.query_execution_ms).toBeGreaterThan(10); 645 | expect(performance.server_total_ms).toBeGreaterThan(10); 646 | expect(performance.response_time).toBeGreaterThan(10); 647 | expect(performance.body_read_time).toBeGreaterThan(0); 648 | expect(performance.compress_time).toBeGreaterThan(0); 649 | expect(performance.deserialize_time).toBeGreaterThan(0); 650 | if (isRuntimeFullyNodeCompatible) { 651 | expect(performance.decompress_time).toEqual(0); // response was too small to compress 652 | } else { 653 | expect(performance.decompress_time).toBeNull; 654 | } 655 | 656 | const billing = resultsWithPerformance.billing; 657 | expect(billing).toEqual({ 658 | billable_logical_bytes_queried: 256000000, 659 | billable_logical_bytes_returned: 24, 660 | }); 661 | 662 | const results2 = await ns.query({ 663 | rank_by: ["vector", "ANN", [1, 1]], 664 | filters: [ 665 | "And", 666 | [ 667 | [ 668 | "Or", 669 | [ 670 | ["numbers", "In", [2, 3]], 671 | ["numbers", "In", [1, 7]], 672 | ], 673 | ], 674 | [ 675 | "Or", 676 | [ 677 | ["foo", "Eq", "bar"], 678 | ["numbers", "In", 4], 679 | ], 680 | ], 681 | ["foo", "NotEq", null], 682 | ["maybeNull", "Eq", null], 683 | ["bool", "Eq", true], 684 | ], 685 | ], 686 | top_k: 10, 687 | }); 688 | expect(results2.rows.length).toEqual(2); 689 | expect(results2.rows[0].id).toEqual(2); 690 | expect(results2.rows[1].id).toEqual(1); 691 | 692 | const recall = await ns.recall({ 693 | num: 1, 694 | top_k: 2, 695 | }); 696 | expect(recall.avg_recall).toEqual(1); 697 | expect(recall.avg_exhaustive_count).toEqual(2); 698 | expect(recall.avg_ann_count).toEqual(2); 699 | 700 | // Delete the second vector. 701 | await ns.write({ 702 | deletes: [1], 703 | }); 704 | 705 | // If we query now, we should only get one result. 706 | const results = await ns.query({ 707 | rank_by: ["vector", "ANN", [1, 1]], 708 | filters: ["numbers", "In", [2, 4]], 709 | top_k: 10, 710 | }); 711 | expect(results.rows.length).toEqual(1); 712 | expect(results.rows[0].id).toEqual(2); 713 | 714 | let num = await ns.approxNumVectors(); 715 | expect(num).toEqual(2); 716 | 717 | let metadata = await ns.metadata(); 718 | expect(metadata.approx_count).toEqual(2); 719 | expect(metadata.dimensions).toEqual(2); 720 | 721 | // Check that `metadata.created_at` data is a valid Date for today, but don't bother checking the 722 | // time. We know it was created today as the test deletes the namespace in the 723 | // beginning. When we compare against the current time, ensure it's UTC. 724 | const now = new Date(); 725 | expect(metadata.created_at.getUTCFullYear()).toEqual(now.getUTCFullYear()); 726 | expect(metadata.created_at.getUTCMonth()).toEqual(now.getUTCMonth()); 727 | expect(metadata.created_at.getUTCDate()).toEqual(now.getUTCDate()); 728 | 729 | // Delete the entire namespace. 730 | await ns.deleteAll(); 731 | 732 | // For some reason, expect().toThrow doesn't catch properly 733 | let gotError: any = null; 734 | try { 735 | await ns.query({ 736 | rank_by: ["vector", "ANN", [1, 1]], 737 | filters: ["numbers", "In", [2, 4]], 738 | top_k: 10, 739 | }); 740 | } catch (e: any) { 741 | gotError = e; 742 | } 743 | expect(gotError).toStrictEqual( 744 | new TurbopufferError(`🤷 namespace '${nameSpaceName}' was not found`, { 745 | status: 404, 746 | }), 747 | ); 748 | }, 10_000); 749 | 750 | test("exists", async () => { 751 | let ns = tpuf.namespace(testNamespacePrefix + "exists"); 752 | 753 | try { 754 | await ns.deleteAll(); 755 | } catch (_: unknown) { 756 | /* empty */ 757 | } 758 | 759 | await ns.write({ 760 | upsert_columns: { 761 | id: [1], 762 | vector: [[0.1, 0.1]], 763 | private: [true], 764 | tags: [["a", "b"]], 765 | }, 766 | distance_metric: "cosine_distance", 767 | }); 768 | 769 | let exists = await ns.exists(); 770 | expect(exists).toEqual(true); 771 | await ns.deleteAll(); 772 | 773 | ns = tpuf.namespace("non_existent_ns"); 774 | exists = await ns.exists(); 775 | expect(exists).toEqual(false); 776 | }); 777 | 778 | const t = isRuntimeFullyNodeCompatible ? it : it.skip; 779 | t("connection_errors_are_wrapped", async () => { 780 | const tpuf = new Turbopuffer({ 781 | baseUrl: "https://api.turbopuffer.com:12345", 782 | apiKey: process.env.TURBOPUFFER_API_KEY!, 783 | connectTimeout: 500, 784 | }); 785 | 786 | const ns = tpuf.namespace( 787 | testNamespacePrefix + "connection_errors_are_wrapped", 788 | ); 789 | 790 | let gotError: any = null; 791 | try { 792 | await ns.query({ 793 | rank_by: ["vector", "ANN", [1, 1]], 794 | top_k: 10, 795 | }); 796 | } catch (e: any) { 797 | gotError = e; 798 | } 799 | 800 | expect(gotError).toStrictEqual( 801 | new TurbopufferError("fetch failed: Connect Timeout Error", {}), 802 | ); 803 | }); 804 | 805 | test("empty_namespace", async () => { 806 | const tpuf = new Turbopuffer({ 807 | apiKey: process.env.TURBOPUFFER_API_KEY!, 808 | }); 809 | 810 | const ns = tpuf.namespace(testNamespacePrefix + "empty_namespace"); 811 | 812 | await ns.write({ 813 | upsert_rows: [ 814 | { 815 | id: 1, 816 | vector: [0.1, 0.1], 817 | }, 818 | ], 819 | distance_metric: "cosine_distance", 820 | }); 821 | 822 | await ns.write({ 823 | deletes: [1], 824 | }); 825 | 826 | await ns.export(); 827 | }); 828 | 829 | test("export", async () => { 830 | const ns = tpuf.namespace(testNamespacePrefix + "export"); 831 | 832 | await ns.write({ 833 | upsert_rows: [ 834 | { 835 | id: 1, 836 | vector: [0.1, 0.1], 837 | title: "one", 838 | private: false, 839 | }, 840 | { 841 | id: 2, 842 | vector: [0.2, 0.2], 843 | title: "two", 844 | private: true, 845 | }, 846 | ], 847 | schema: { 848 | title: { 849 | type: "string", 850 | full_text_search: true, 851 | }, 852 | private: { 853 | type: "bool", 854 | }, 855 | }, 856 | distance_metric: "cosine_distance", 857 | }); 858 | 859 | const data = await ns.export(); 860 | expect(data).toEqual({ 861 | ids: [1, 2], 862 | vectors: [ 863 | [0.1, 0.1], 864 | [0.2, 0.2], 865 | ], 866 | attributes: { 867 | private: [false, true], 868 | title: ["one", "two"], 869 | }, 870 | next_cursor: null, 871 | }); 872 | 873 | await ns.deleteAll(); 874 | }); 875 | 876 | test("no_cmek", async () => { 877 | const ns = tpuf.namespace(testNamespacePrefix + "no_cmek"); 878 | 879 | let error: any = null; 880 | try { 881 | await ns.write({ 882 | upsert_rows: [ 883 | { 884 | id: 1, 885 | vector: [0.1, 0.1], 886 | }, 887 | ], 888 | distance_metric: "cosine_distance", 889 | encryption: { 890 | cmek: { 891 | key_name: "mykey", 892 | }, 893 | }, 894 | }); 895 | } catch (e: any) { 896 | error = e; 897 | } 898 | 899 | expect(error).toBeInstanceOf(TurbopufferError); 900 | }); 901 | 902 | test("copy_from_namespace", async () => { 903 | const ns1Name = testNamespacePrefix + "copy_from_namespace_1"; 904 | const ns1 = tpuf.namespace(ns1Name); 905 | const ns2 = tpuf.namespace(testNamespacePrefix + "copy_from_namespace_2"); 906 | 907 | try { 908 | await ns1.deleteAll(); 909 | await ns2.deleteAll(); 910 | } catch (_: unknown) { 911 | /* empty */ 912 | } 913 | 914 | // let's test with a columnar write 915 | await ns1.write({ 916 | upsert_columns: { 917 | id: [1, 2, 3], 918 | vector: [ 919 | [0.1, 0.1], 920 | [0.2, 0.2], 921 | [0.3, 0.3], 922 | ], 923 | tags: [["a"], ["b"], ["c"]], 924 | }, 925 | distance_metric: "cosine_distance", 926 | }); 927 | 928 | await ns2.copyFromNamespace(ns1Name); 929 | 930 | const res = await ns2.query({ 931 | rank_by: ["vector", "ANN", [0.1, 0.1]], 932 | include_attributes: true, 933 | top_k: 10, 934 | }); 935 | 936 | expect(res.rows.length).toEqual(3); 937 | }); 938 | 939 | test("patch", async () => { 940 | const ns = tpuf.namespace(testNamespacePrefix + "patch"); 941 | 942 | try { 943 | await ns.deleteAll(); 944 | } catch (_: unknown) { 945 | /* empty */ 946 | } 947 | 948 | await ns.write({ 949 | upsert_rows: [ 950 | { 951 | id: 1, 952 | vector: [1, 1], 953 | }, 954 | { 955 | id: 2, 956 | vector: [2, 2], 957 | }, 958 | ], 959 | distance_metric: "cosine_distance", 960 | }); 961 | 962 | await ns.write({ 963 | patch_rows: [ 964 | { id: 1, a: 1 }, 965 | { id: 2, b: 2 }, 966 | ], 967 | }); 968 | 969 | await ns.write({ 970 | patch_rows: [ 971 | { id: 1, b: 1 }, 972 | { id: 2, a: 2 }, 973 | ], 974 | }); 975 | 976 | let results = await ns.query({ rank_by: ["id", "asc"], include_attributes: ['id', 'a', 'b'], top_k: 10 }); 977 | expect(results.rows.length).toEqual(2); 978 | expect(results.rows[0]).toEqual({ id: 1, a: 1, b: 1 }); 979 | expect(results.rows[1]).toEqual({ id: 2, a: 2, b: 2 }); 980 | 981 | await ns.write({ 982 | patch_columns: { 983 | id: [1, 2], 984 | a: [11, 22], 985 | c: [1, 2], 986 | }, 987 | }); 988 | 989 | results = await ns.query({ rank_by: ["id", "asc"], include_attributes: ['id', 'a', 'b', 'c'], top_k: 10 }); 990 | expect(results.rows.length).toEqual(2); 991 | expect(results.rows[0]).toEqual({ id: 1, a: 11, b: 1, c: 1 }); 992 | expect(results.rows[1]).toEqual({ id: 2, a: 22, b: 2, c: 2 }); 993 | 994 | await ns.deleteAll(); 995 | }); 996 | 997 | test("delete_by_filter", async () => { 998 | const ns = tpuf.namespace(testNamespacePrefix + "delete_by_filter"); 999 | 1000 | try { 1001 | await ns.deleteAll(); 1002 | } catch (_: unknown) { 1003 | /* empty */ 1004 | } 1005 | 1006 | await ns.write({ 1007 | upsert_rows: [ 1008 | { 1009 | id: 1, 1010 | vector: [1, 2], 1011 | foo: "bar", 1012 | }, 1013 | { 1014 | id: 2, 1015 | vector: [3, 4], 1016 | foo: "baz", 1017 | }, 1018 | { 1019 | id: 3, 1020 | vector: [3, 4], 1021 | foo: "baz", 1022 | }, 1023 | ], 1024 | distance_metric: "cosine_distance", 1025 | }); 1026 | 1027 | let results = await ns.query({ rank_by: ["id", "asc"], top_k: 10 }); 1028 | expect(results.rows.length).toEqual(3); 1029 | 1030 | const rowsAffected = await ns.write({ 1031 | delete_by_filter: ["foo", "Eq", "baz"], 1032 | }); 1033 | expect(rowsAffected).toEqual(2); 1034 | 1035 | results = await ns.query({ rank_by: ["id", "asc"], top_k: 10 }); 1036 | expect(results.rows.length).toEqual(1); 1037 | expect(results.rows[0].id).toEqual(1); 1038 | 1039 | await ns.deleteAll(); 1040 | }); 1041 | 1042 | function randomVector(dims: number) { 1043 | return Array(dims) 1044 | .fill(0) 1045 | .map(() => Math.random()); 1046 | } 1047 | 1048 | test("compression", async () => { 1049 | const ns = tpuf.namespace(testNamespacePrefix + "compression"); 1050 | 1051 | try { 1052 | await ns.deleteAll(); 1053 | } catch (_: unknown) { 1054 | /* empty */ 1055 | } 1056 | 1057 | // Insert a large number of vectors to trigger compression 1058 | await ns.write({ 1059 | upsert_rows: Array.from({ length: 10 }, (_, i) => ({ 1060 | id: i + 1, 1061 | vector: randomVector(1024), 1062 | text: "b".repeat(1024), 1063 | })), 1064 | distance_metric: "cosine_distance", 1065 | }); 1066 | 1067 | const resultsWithPerformance = await ns.query({ 1068 | rank_by: ["vector", "ANN", randomVector(1024)], 1069 | top_k: 10, 1070 | include_attributes: true, 1071 | }); 1072 | 1073 | const performance = resultsWithPerformance.performance; 1074 | expect(performance.compress_time).toBeGreaterThan(0); 1075 | expect(performance.body_read_time).toBeGreaterThan(0); 1076 | expect(performance.deserialize_time).toBeGreaterThan(0); 1077 | if (isRuntimeFullyNodeCompatible) { 1078 | expect(performance.decompress_time).toBeGreaterThan(0); // Response should be compressed 1079 | } else { 1080 | expect(performance.decompress_time).toBeNull; 1081 | } 1082 | }); 1083 | 1084 | test("disable_compression", async () => { 1085 | const tpufNoCompression = new Turbopuffer({ 1086 | apiKey: process.env.TURBOPUFFER_API_KEY!, 1087 | compression: false, 1088 | }); 1089 | 1090 | const ns = tpufNoCompression.namespace( 1091 | testNamespacePrefix + "disable_compression", 1092 | ); 1093 | 1094 | try { 1095 | await ns.deleteAll(); 1096 | } catch (_: unknown) { 1097 | /* empty */ 1098 | } 1099 | 1100 | // Insert a large number of vectors to trigger compression 1101 | await ns.write({ 1102 | upsert_rows: Array.from({ length: 10 }, (_, i) => ({ 1103 | id: i + 1, 1104 | vector: randomVector(1024), 1105 | text: "b".repeat(1024), 1106 | })), 1107 | distance_metric: "cosine_distance", 1108 | }); 1109 | 1110 | const resultsWithPerformance = await ns.query({ 1111 | rank_by: ["vector", "ANN", randomVector(1024)], 1112 | top_k: 10, 1113 | include_attributes: true, 1114 | }); 1115 | 1116 | const performance = resultsWithPerformance.performance; 1117 | expect(performance.compress_time).toBeNull; 1118 | expect(performance.body_read_time).toBeGreaterThan(0); 1119 | expect(performance.deserialize_time).toBeGreaterThan(0); 1120 | if (isRuntimeFullyNodeCompatible) { 1121 | expect(performance.decompress_time).toEqual(0); 1122 | } else { 1123 | expect(performance.decompress_time).toBeNull; 1124 | } 1125 | }); 1126 | 1127 | test("product_operator", async () => { 1128 | const ns = tpuf.namespace(testNamespacePrefix + "product_operator"); 1129 | 1130 | try { 1131 | await ns.deleteAll(); 1132 | } catch (_: unknown) { 1133 | /* empty */ 1134 | } 1135 | 1136 | const schema: Schema = { 1137 | title: { 1138 | type: "string", 1139 | full_text_search: true, 1140 | }, 1141 | content: { 1142 | type: "string", 1143 | full_text_search: true, 1144 | }, 1145 | }; 1146 | 1147 | await ns.write({ 1148 | upsert_rows: [ 1149 | { 1150 | id: 1, 1151 | vector: [0.1, 0.1], 1152 | title: "one", 1153 | content: "foo bar baz", 1154 | }, 1155 | { 1156 | id: 2, 1157 | vector: [0.2, 0.2], 1158 | title: "two", 1159 | content: "foo bar", 1160 | }, 1161 | { 1162 | id: 3, 1163 | vector: [0.3, 0.3], 1164 | title: "three", 1165 | content: "bar baz", 1166 | }, 1167 | ], 1168 | distance_metric: "euclidean_squared", 1169 | schema: schema, 1170 | }); 1171 | 1172 | const queries: RankBy[] = [ 1173 | ["Product", [2, ["title", "BM25", "one"]]], 1174 | ["Product", [["title", "BM25", "one"], 2]], 1175 | [ 1176 | "Sum", 1177 | [ 1178 | ["Product", [2, ["title", "BM25", "one"]]], 1179 | ["content", "BM25", "foo"], 1180 | ], 1181 | ], 1182 | [ 1183 | "Product", 1184 | [ 1185 | 2, 1186 | [ 1187 | "Max", 1188 | [ 1189 | ["Product", [2, ["title", "BM25", "one"]]], 1190 | ["content", "BM25", "foo"], 1191 | ], 1192 | ], 1193 | ], 1194 | ], 1195 | ]; 1196 | 1197 | for (const query of queries) { 1198 | const results = await ns.query({ rank_by: query, top_k: 10 }); 1199 | expect(results.rows.length).toBeGreaterThan(0); 1200 | } 1201 | }); 1202 | 1203 | test("not", async () => { 1204 | const ns = tpuf.namespace(testNamespacePrefix + "not"); 1205 | try { 1206 | await ns.deleteAll(); 1207 | } catch (_: unknown) { 1208 | /* empty */ 1209 | } 1210 | 1211 | await ns.write({ 1212 | upsert_columns: { 1213 | id: [1], 1214 | vector: [[0.1, 0.1]], 1215 | text: ["Walruses are large marine mammals with long tusks and whiskers"], 1216 | }, 1217 | schema: { 1218 | text: { 1219 | type: "string", 1220 | full_text_search: { 1221 | stemming: true, 1222 | }, 1223 | }, 1224 | }, 1225 | distance_metric: "cosine_distance", 1226 | }); 1227 | 1228 | const results = await ns.query({ 1229 | rank_by: ["text", "BM25", "walrus whisker"], 1230 | filters: ["text", "ContainsAllTokens", "marine mammals"], 1231 | top_k: 10, 1232 | }); 1233 | expect(results.rows.length).toEqual(1); 1234 | 1235 | const resultsNot0 = await ns.query({ 1236 | rank_by: ["text", "BM25", "walrus whisker"], 1237 | filters: ["Not", ["text", "ContainsAllTokens", "marine mammals"]], 1238 | top_k: 10, 1239 | }); 1240 | expect(resultsNot0.rows.length).toEqual(0); 1241 | 1242 | const resultsNot1 = await ns.query({ 1243 | rank_by: ["text", "BM25", "walrus whisker"], 1244 | filters: ["Not", ["Not", ["text", "ContainsAllTokens", "marine mammals"]]], 1245 | top_k: 10, 1246 | }); 1247 | expect(resultsNot1.rows.length).toEqual(1); 1248 | 1249 | const resultsNot2 = await ns.query({ 1250 | rank_by: ["text", "BM25", "walrus whisker"], 1251 | filters: [ 1252 | "Or", 1253 | [ 1254 | ["text", "ContainsAllTokens", "marine things"], 1255 | [ 1256 | "Or", 1257 | [ 1258 | [ 1259 | "Not", 1260 | [ 1261 | "Not", 1262 | [ 1263 | "And", 1264 | [ 1265 | ["text", "ContainsAllTokens", "marine mammals"], 1266 | ["id", "In", [0, 1, 2]], 1267 | ], 1268 | ], 1269 | ], 1270 | ], 1271 | ], 1272 | ], 1273 | ], 1274 | ], 1275 | top_k: 10, 1276 | }); 1277 | expect(resultsNot2.rows.length).toEqual(1); 1278 | }); 1279 | 1280 | test("readme", async () => { 1281 | const ns = tpuf.namespace(testNamespacePrefix + "readme"); 1282 | 1283 | try { 1284 | await ns.deleteAll(); 1285 | } catch (_: unknown) { 1286 | /* empty */ 1287 | } 1288 | 1289 | await ns.write({ 1290 | upsert_rows: [ 1291 | { 1292 | id: 1, 1293 | vector: [1, 2], 1294 | foo: "bar", 1295 | numbers: [1, 2, 3], 1296 | }, 1297 | { 1298 | id: 2, 1299 | vector: [3, 4], 1300 | foo: "baz", 1301 | numbers: [2, 3, 4], 1302 | }, 1303 | ], 1304 | distance_metric: "cosine_distance", 1305 | }); 1306 | 1307 | const results = await ns.query({ 1308 | rank_by: ["vector", "ANN", [1, 1]], 1309 | filters: ["numbers", "In", [2, 4]], 1310 | top_k: 10, 1311 | }); 1312 | 1313 | expect(results.rows.length).toEqual(2); 1314 | expect(results.rows[0].id).toEqual(2); 1315 | expect(results.rows[0].$dist).toBeGreaterThanOrEqual(0); 1316 | expect(results.rows[1].id).toEqual(1); 1317 | expect(results.rows[1].$dist).toBeGreaterThanOrEqual(0); 1318 | 1319 | await ns.deleteAll(); 1320 | }); 1321 | 1322 | // test helper and utility methods 1323 | 1324 | test("test_buildBaseUrl", () => { 1325 | // if no protocol, add https:// 1326 | expect(buildBaseUrl("gcp-us-east4.turbopuffer.com")).toEqual( 1327 | "https://gcp-us-east4.turbopuffer.com", 1328 | ); 1329 | 1330 | // if any protocol (or protocol-looking string) exists, do nothing 1331 | expect(buildBaseUrl("https://gcp-us-east4.turbopuffer.com")).toEqual( 1332 | "https://gcp-us-east4.turbopuffer.com", 1333 | ); 1334 | expect(buildBaseUrl("http://gcp-us-east4.turbopuffer.com")).toEqual( 1335 | "http://gcp-us-east4.turbopuffer.com", 1336 | ); 1337 | expect(buildBaseUrl("admin://gcp-us-east4.turbopuffer.com")).toEqual( 1338 | "admin://gcp-us-east4.turbopuffer.com", 1339 | ); 1340 | 1341 | // do not add trailing / 1342 | expect(buildBaseUrl("https://gcp-us-east4.turbopuffer.com")).not.toEqual( 1343 | "https://gcp-us-east4.turbopuffer.com/", 1344 | ); 1345 | }); 1346 | 1347 | test("test_buildUrl", () => { 1348 | /** baseUrl w/o path **/ 1349 | /* w/o path + w/o query */ 1350 | expect(buildUrl("https://gcp-us-east4.turbopuffer.com", "").href).toEqual( 1351 | "https://gcp-us-east4.turbopuffer.com/", 1352 | ); 1353 | 1354 | /* w/o path + w/ query */ 1355 | expect( 1356 | buildUrl("https://gcp-us-east4.turbopuffer.com", "/v1/namespaces", { 1357 | cursor: "next_cursor", 1358 | prefix: "my_prefix", 1359 | page_size: "15", 1360 | }).href, 1361 | ).toEqual( 1362 | "https://gcp-us-east4.turbopuffer.com/v1/namespaces?cursor=next_cursor&prefix=my_prefix&page_size=15", 1363 | ); 1364 | 1365 | /* w/ path + w/o query */ 1366 | expect( 1367 | buildUrl("https://gcp-us-east4.turbopuffer.com", "/v1/namespaces").href, 1368 | ).toEqual("https://gcp-us-east4.turbopuffer.com/v1/namespaces"); 1369 | 1370 | expect( 1371 | buildUrl("https://gcp-us-east4.turbopuffer.com", "v1/namespaces").href, 1372 | ).toEqual("https://gcp-us-east4.turbopuffer.com/v1/namespaces"); 1373 | 1374 | expect( 1375 | buildUrl("https://gcp-us-east4.turbopuffer.com", "v1/namespaces/").href, 1376 | ).toEqual("https://gcp-us-east4.turbopuffer.com/v1/namespaces/"); 1377 | 1378 | expect( 1379 | buildUrl("https://gcp-us-east4.turbopuffer.com/", "/v1/namespaces").href, 1380 | ).toEqual("https://gcp-us-east4.turbopuffer.com/v1/namespaces"); 1381 | 1382 | expect( 1383 | buildUrl("https://gcp-us-east4.turbopuffer.com/", "v1/namespaces").href, 1384 | ).toEqual("https://gcp-us-east4.turbopuffer.com/v1/namespaces"); 1385 | 1386 | expect( 1387 | buildUrl("https://gcp-us-east4.turbopuffer.com//", "/v1/namespaces").href, 1388 | ).toEqual("https://gcp-us-east4.turbopuffer.com/v1/namespaces"); 1389 | 1390 | expect( 1391 | buildUrl("https://gcp-us-east4.turbopuffer.com//", "//v1/namespaces").href, 1392 | ).toEqual("https://gcp-us-east4.turbopuffer.com/v1/namespaces"); 1393 | 1394 | /* w/ path + w/ query */ 1395 | expect( 1396 | buildUrl("https://gcp-us-east4.turbopuffer.com", "/v1/namespaces", { 1397 | cursor: "next_cursor", 1398 | prefix: "my_prefix", 1399 | page_size: "15", 1400 | }).href, 1401 | ).toEqual( 1402 | "https://gcp-us-east4.turbopuffer.com/v1/namespaces?cursor=next_cursor&prefix=my_prefix&page_size=15", 1403 | ); 1404 | 1405 | /** baseUrl w/ path **/ 1406 | /* w/o path + w/o query */ 1407 | expect( 1408 | buildUrl("https://gcp-us-east4.turbopuffer.com/my-cool-path", "").href, 1409 | ).toEqual("https://gcp-us-east4.turbopuffer.com/my-cool-path/"); 1410 | 1411 | /* w/o path + w/ query */ 1412 | expect( 1413 | buildUrl("https://gcp-us-east4.turbopuffer.com/my-cool-path", "", { 1414 | cursor: "next_cursor", 1415 | prefix: "my_prefix", 1416 | page_size: "15", 1417 | }).href, 1418 | ).toEqual( 1419 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/?cursor=next_cursor&prefix=my_prefix&page_size=15", 1420 | ); 1421 | 1422 | expect( 1423 | buildUrl("https://gcp-us-east4.turbopuffer.com/my-cool-path/", "", { 1424 | cursor: "next_cursor", 1425 | prefix: "my_prefix", 1426 | page_size: "15", 1427 | }).href, 1428 | ).toEqual( 1429 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/?cursor=next_cursor&prefix=my_prefix&page_size=15", 1430 | ); 1431 | 1432 | /* w/ path + w/o query */ 1433 | expect( 1434 | buildUrl( 1435 | "https://gcp-us-east4.turbopuffer.com/my-cool-path", 1436 | "/v1/namespaces", 1437 | ).href, 1438 | ).toEqual("https://gcp-us-east4.turbopuffer.com/my-cool-path/v1/namespaces"); 1439 | 1440 | expect( 1441 | buildUrl( 1442 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/", 1443 | "/v1/namespaces", 1444 | ).href, 1445 | ).toEqual("https://gcp-us-east4.turbopuffer.com/my-cool-path/v1/namespaces"); 1446 | 1447 | expect( 1448 | buildUrl( 1449 | "https://gcp-us-east4.turbopuffer.com/my-cool-path//", 1450 | "/v1/namespaces", 1451 | ).href, 1452 | ).toEqual("https://gcp-us-east4.turbopuffer.com/my-cool-path/v1/namespaces"); 1453 | 1454 | expect( 1455 | buildUrl( 1456 | "https://gcp-us-east4.turbopuffer.com/my-cool-path//", 1457 | "v1/namespaces", 1458 | ).href, 1459 | ).toEqual("https://gcp-us-east4.turbopuffer.com/my-cool-path/v1/namespaces"); 1460 | 1461 | /* w/ path + w/ query */ 1462 | expect( 1463 | buildUrl( 1464 | "https://gcp-us-east4.turbopuffer.com/my-cool-path//", 1465 | "v1/namespaces", 1466 | { 1467 | cursor: "next_cursor", 1468 | prefix: "my_prefix", 1469 | page_size: "15", 1470 | }, 1471 | ).href, 1472 | ).toEqual( 1473 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/v1/namespaces?cursor=next_cursor&prefix=my_prefix&page_size=15", 1474 | ); 1475 | 1476 | expect( 1477 | buildUrl( 1478 | "https://gcp-us-east4.turbopuffer.com/my-cool-path//", 1479 | "v1/namespaces", 1480 | { 1481 | cursor: "next_cursor", 1482 | prefix: "my_prefix", 1483 | page_size: "15", 1484 | }, 1485 | ).href, 1486 | ).toEqual( 1487 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/v1/namespaces?cursor=next_cursor&prefix=my_prefix&page_size=15", 1488 | ); 1489 | 1490 | /** baseUrl w/ double path **/ 1491 | expect( 1492 | buildUrl( 1493 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/another-dope-path", 1494 | "", 1495 | ).href, 1496 | ).toEqual( 1497 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/another-dope-path/", 1498 | ); 1499 | 1500 | expect( 1501 | buildUrl( 1502 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/another-dope-path", 1503 | "", 1504 | { 1505 | cursor: "next_cursor", 1506 | prefix: "my_prefix", 1507 | page_size: "15", 1508 | }, 1509 | ).href, 1510 | ).toEqual( 1511 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/another-dope-path/?cursor=next_cursor&prefix=my_prefix&page_size=15", 1512 | ); 1513 | 1514 | expect( 1515 | buildUrl( 1516 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/another-dope-path", 1517 | "/v1/namespaces", 1518 | ).href, 1519 | ).toEqual( 1520 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/another-dope-path/v1/namespaces", 1521 | ); 1522 | 1523 | expect( 1524 | buildUrl( 1525 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/another-dope-path", 1526 | "/v1/namespaces", 1527 | { 1528 | cursor: "next_cursor", 1529 | prefix: "my_prefix", 1530 | page_size: "15", 1531 | }, 1532 | ).href, 1533 | ).toEqual( 1534 | "https://gcp-us-east4.turbopuffer.com/my-cool-path/another-dope-path/v1/namespaces?cursor=next_cursor&prefix=my_prefix&page_size=15", 1535 | ); 1536 | }); 1537 | -------------------------------------------------------------------------------- /src/turbopuffer.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * Official TypeScript SDK for turbopuffer.com's API 3 | * Something missing or should be improved? Email morgan@turbopuffer.com. 4 | * 5 | * Based off the initial work of https://github.com/holocron-hq! Thank you ❤️ 6 | */ 7 | 8 | import { createHTTPClient } from "./createHTTPClient"; 9 | import { 10 | shouldCompressWrite, 11 | TurbopufferError, 12 | } from "./helpers"; 13 | import type { 14 | Consistency, 15 | ExportResponse, 16 | Filters, 17 | HintCacheWarmResponse, 18 | HTTPClient, 19 | NamespaceMetadata, 20 | NamespacesListResult, 21 | QueryResults, 22 | RankBy, 23 | RecallMeasurement, 24 | Schema, 25 | WriteParams, 26 | } from "./types"; 27 | 28 | /* Base Client */ 29 | export class Turbopuffer { 30 | http: HTTPClient; 31 | 32 | constructor({ 33 | apiKey, 34 | baseUrl = "https://api.turbopuffer.com", 35 | connectTimeout = 10 * 1000, 36 | connectionIdleTimeout = 60 * 1000, 37 | warmConnections = 0, 38 | compression = true, 39 | }: { 40 | /** The API key to authenticate with. */ 41 | apiKey: string; 42 | /** The base URL. Default is https://api.turbopuffer.com. */ 43 | baseUrl?: string; 44 | /** The timeout to establish a connection, in ms. Default is 10_000. Only applicable in Node and Deno.*/ 45 | connectTimeout?: number; 46 | /** The socket idle timeout, in ms. Default is 60_000. Only applicable in Node and Deno.*/ 47 | connectionIdleTimeout?: number; 48 | /** The number of connections to open initially when creating a new client. Default is 0. */ 49 | warmConnections?: number; 50 | /** Whether to compress requests and accept compressed responses. Default is true. */ 51 | compression?: boolean; 52 | }) { 53 | this.http = createHTTPClient( 54 | baseUrl, 55 | apiKey, 56 | connectTimeout, 57 | connectionIdleTimeout, 58 | warmConnections, 59 | compression, 60 | ); 61 | } 62 | 63 | /** 64 | * List all your namespaces. 65 | * See: https://turbopuffer.com/docs/namespaces 66 | */ 67 | async namespaces({ 68 | cursor, 69 | prefix, 70 | page_size, 71 | }: { 72 | cursor?: string; 73 | prefix?: string; 74 | page_size?: number; 75 | } = {}): Promise { 76 | return ( 77 | await this.http.doRequest({ 78 | method: "GET", 79 | path: "/v1/namespaces", 80 | query: { 81 | cursor, 82 | prefix, 83 | page_size: page_size ? page_size.toString() : undefined, 84 | }, 85 | retryable: true, 86 | }) 87 | ).body!; 88 | } 89 | 90 | /** 91 | * Creates a namespace object to operate on. Operations 92 | * should be called on the Namespace object itself. 93 | */ 94 | namespace(id: string): Namespace { 95 | return new Namespace(this, id); 96 | } 97 | } 98 | 99 | export class Namespace { 100 | private client: Turbopuffer; 101 | id: string; 102 | 103 | constructor(client: Turbopuffer, id: string) { 104 | this.client = client; 105 | this.id = id; 106 | } 107 | 108 | async write(params: WriteParams): Promise { 109 | const response = await this.client.http.doRequest<{ 110 | status: string; 111 | rows_affected: number; 112 | }>({ 113 | method: "POST", 114 | path: `/v2/namespaces/${this.id}`, 115 | compress: shouldCompressWrite(params), 116 | body: params, 117 | retryable: true, // writes are idempotent 118 | }); 119 | return response.body?.rows_affected ?? 0; 120 | } 121 | 122 | /** 123 | * Queries vectors. 124 | * See: https://turbopuffer.com/docs/query 125 | */ 126 | async query({ 127 | ...params 128 | }: { 129 | top_k: number; 130 | include_attributes?: boolean | string[]; 131 | filters?: Filters; 132 | rank_by: RankBy; 133 | consistency?: Consistency; 134 | }): Promise { 135 | const response = await this.client.http.doRequest({ 136 | method: "POST", 137 | path: `/v2/namespaces/${this.id}/query`, 138 | body: params, 139 | retryable: true, 140 | compress: true, 141 | }); 142 | 143 | const results = response.body!; 144 | results.performance = { 145 | ...results.performance, 146 | ...response.request_timing, 147 | }; 148 | 149 | return results; 150 | } 151 | 152 | /** 153 | * Warm the cache. 154 | */ 155 | async hintCacheWarm(): Promise { 156 | return ( 157 | await this.client.http.doRequest({ 158 | method: "GET", 159 | path: `/v1/namespaces/${this.id}/hint_cache_warm`, 160 | retryable: true, 161 | }) 162 | ).body!; 163 | } 164 | 165 | /** 166 | * Export all vectors at full precision. 167 | * See: https://turbopuffer.com/docs/export 168 | */ 169 | async export(params?: { cursor?: string }): Promise { 170 | const response = await this.client.http.doRequest({ 171 | method: "GET", 172 | path: `/v1/namespaces/${this.id}`, 173 | query: { cursor: params?.cursor }, 174 | retryable: true, 175 | }); 176 | const { ids, vectors, attributes, next_cursor } = response.body!; 177 | return { 178 | ids, 179 | vectors, 180 | attributes, 181 | next_cursor, 182 | }; 183 | } 184 | 185 | /** 186 | * Fetches the approximate number of vectors in a namespace. 187 | */ 188 | async approxNumVectors(): Promise { 189 | return (await this.metadata()).approx_count; 190 | } 191 | 192 | async metadata(): Promise { 193 | const response = await this.client.http.doRequest({ 194 | method: "HEAD", 195 | path: `/v1/namespaces/${this.id}`, 196 | retryable: true, 197 | }); 198 | 199 | return { 200 | id: this.id, 201 | approx_count: parseInt( 202 | response.headers["x-turbopuffer-approx-num-vectors"], 203 | ), 204 | dimensions: parseInt(response.headers["x-turbopuffer-dimensions"]), 205 | created_at: new Date(response.headers["x-turbopuffer-created-at"]), 206 | }; 207 | } 208 | 209 | /** 210 | * Checks if a namespace exists. 211 | */ 212 | async exists(): Promise { 213 | try { 214 | await this.metadata(); 215 | return true; 216 | } catch (e) { 217 | if (e instanceof TurbopufferError && e.status === 404) { 218 | return false; 219 | } 220 | throw e; 221 | } 222 | } 223 | 224 | /** 225 | * Delete a namespace fully (all data). 226 | * See: https://turbopuffer.com/docs/delete-namespace 227 | */ 228 | async deleteAll(): Promise { 229 | await this.client.http.doRequest<{ status: string }>({ 230 | method: "DELETE", 231 | path: `/v1/namespaces/${this.id}`, 232 | retryable: true, 233 | }); 234 | } 235 | 236 | /** 237 | * Evaluates the recall performance of ANN queries in a namespace. 238 | * See: https://turbopuffer.com/docs/recall 239 | */ 240 | async recall({ 241 | num, 242 | top_k, 243 | filters, 244 | queries, 245 | }: { 246 | num?: number; 247 | top_k?: number; 248 | filters?: Filters; 249 | queries?: number[][]; 250 | }): Promise { 251 | return ( 252 | await this.client.http.doRequest({ 253 | method: "POST", 254 | path: `/v1/namespaces/${this.id}/_debug/recall`, 255 | compress: queries && queries.length > 10, 256 | body: { 257 | num, 258 | top_k, 259 | filters, 260 | queries: queries 261 | ? queries.reduce((acc, value) => acc.concat(value), []) 262 | : undefined, 263 | }, 264 | retryable: true, 265 | }) 266 | ).body!; 267 | } 268 | 269 | /** 270 | * Returns the current schema for the namespace. 271 | * See: https://turbopuffer.com/docs/schema 272 | */ 273 | async schema(): Promise { 274 | return ( 275 | await this.client.http.doRequest({ 276 | method: "GET", 277 | path: `/v1/namespaces/${this.id}/schema`, 278 | retryable: true, 279 | }) 280 | ).body!; 281 | } 282 | 283 | /** 284 | * Updates the schema for a namespace. 285 | * Returns the final schema after updates are done. 286 | * See https://turbopuffer.com/docs/schema for specifics on allowed updates. 287 | */ 288 | async updateSchema(updatedSchema: Schema): Promise { 289 | return ( 290 | await this.client.http.doRequest({ 291 | method: "POST", 292 | path: `/v1/namespaces/${this.id}/schema`, 293 | body: updatedSchema, 294 | retryable: true, 295 | }) 296 | ).body!; 297 | } 298 | 299 | /** 300 | * Copies all documents from another namespace to this namespace. 301 | * See: https://turbopuffer.com/docs/upsert#parameters `copy_from_namespace` 302 | * for specifics on how this works. 303 | */ 304 | async copyFromNamespace(sourceNamespace: string) { 305 | await this.client.http.doRequest({ 306 | method: "POST", 307 | path: `/v1/namespaces/${this.id}`, 308 | body: { 309 | copy_from_namespace: sourceNamespace, 310 | }, 311 | retryable: true, 312 | }); 313 | } 314 | } 315 | -------------------------------------------------------------------------------- /src/types.ts: -------------------------------------------------------------------------------- 1 | // Utility Types 2 | // Note: At the moment, negative numbers aren't supported. 3 | export type Id = string | number; 4 | 5 | export type AttributeType = 6 | | null 7 | | string 8 | | number 9 | | string[] 10 | | number[] 11 | | boolean; 12 | 13 | export interface FTSParams { 14 | k1: number; 15 | b: number; 16 | language: string; 17 | stemming: boolean; 18 | remove_stopwords: boolean; 19 | case_sensitive: boolean; 20 | tokenizer: string; 21 | } 22 | export type SchemaType = 23 | | "string" 24 | | "int" 25 | | "uint" 26 | | "uuid" 27 | | "datetime" 28 | | "bool" 29 | | "[]string" 30 | | "[]int" 31 | | "[]uint" 32 | | "[]uuid" 33 | | `[${number}]${"f16" | "f32"}`; 34 | 35 | export type Schema = Record< 36 | string, 37 | { 38 | type?: SchemaType; 39 | filterable?: boolean; 40 | full_text_search?: boolean | Partial; 41 | ann?: boolean; 42 | } 43 | >; 44 | 45 | export type RankBy_Vector = ["vector", "ANN", number[]]; 46 | export type RankBy_OrderByAttribute = [string, "asc" | "desc"]; 47 | export type RankBy_Text = 48 | | [string, "BM25", string | string[]] 49 | | ["Sum" | "Max", RankBy_Text[]] 50 | | ["Product", [RankBy_Text, number]] 51 | | ["Product", [number, RankBy_Text]]; 52 | export type RankBy = RankBy_Vector | RankBy_Text | RankBy_OrderByAttribute; 53 | 54 | export interface Consistency { 55 | level: "strong" | "eventual"; 56 | } 57 | 58 | export type DistanceMetric = "cosine_distance" | "euclidean_squared"; 59 | export type FilterOperator = 60 | | "Eq" 61 | | "NotEq" 62 | | "In" 63 | | "NotIn" 64 | | "Lt" 65 | | "Lte" 66 | | "Gt" 67 | | "Gte" 68 | | "Glob" 69 | | "NotGlob" 70 | | "IGlob" 71 | | "NotIGlob" 72 | | "ContainsAllTokens"; 73 | export type FilterConnective = "And" | "Or"; 74 | export type Not = "Not"; 75 | export type FilterValue = AttributeType; 76 | export type FilterCondition = [string, FilterOperator, FilterValue]; 77 | export type Filters = 78 | | [FilterConnective, Filters[]] 79 | | [Not, Filters] 80 | | FilterCondition; 81 | export interface Cmek { 82 | key_name: string; 83 | } 84 | export interface Encryption { 85 | cmek: Cmek; 86 | } 87 | 88 | export interface RequestParams { 89 | method: string; 90 | path: string; 91 | query?: Record; 92 | body?: unknown; 93 | compress?: boolean; 94 | retryable?: boolean; 95 | } 96 | 97 | export interface RequestTiming { 98 | response_time: number; 99 | body_read_time: number | null; 100 | decompress_time: number | null; 101 | compress_time: number | null; 102 | deserialize_time: number | null; 103 | } 104 | 105 | export type RequestResponse = Promise<{ 106 | body?: T; 107 | headers: Record; 108 | request_timing: RequestTiming; 109 | }>; 110 | 111 | export interface HTTPClient { 112 | doRequest(_: RequestParams): RequestResponse; 113 | } 114 | 115 | export interface TpufResponseWithMetadata { 116 | body_text: string; 117 | body_read_end: number; 118 | decompress_end: number; 119 | } 120 | 121 | export interface ColumnarDocs { 122 | id: Id[]; 123 | /** 124 | * Required if the namespace has a vector index. 125 | * For non-vector namespaces, this key should be omitted. 126 | */ 127 | vector?: number[][]; 128 | } 129 | export type ColumnarAttributes = Record; 130 | 131 | export type UpsertColumns = ColumnarDocs & ColumnarAttributes; 132 | export type PatchColumns = { id: Id[] } & ColumnarAttributes; 133 | 134 | interface RowDoc { 135 | id: Id; 136 | /** 137 | * Required if the namespace has a vector index. 138 | * For non-vector namespaces, this key should be omitted. 139 | */ 140 | vector?: number[]; 141 | } 142 | type RowAttributes = Record; 143 | 144 | export type UpsertRows = (RowDoc & RowAttributes)[]; 145 | export type PatchRows = ({ id: Id } & RowAttributes)[]; 146 | 147 | export interface WriteParams { 148 | /** Upserts documents in a column-based format. */ 149 | upsert_columns?: UpsertColumns; 150 | /** Upserts documents in a row-based format. */ 151 | upsert_rows?: UpsertRows; 152 | /** 153 | * Patches documents in a column-based format. Identical to `upsert_columns`, but 154 | * instead of overwriting entire documents, only the specified keys are written. 155 | */ 156 | patch_columns?: PatchColumns; 157 | /** 158 | * Patches documents in a row-based format. Identical to `upsert_rows`, but 159 | * instead of overwriting entire documents, only the specified keys are written. 160 | */ 161 | patch_rows?: PatchRows; 162 | /** Deletes documents by ID. */ 163 | deletes?: Id[]; 164 | /** Deletes documents that match a filter. */ 165 | delete_by_filter?: Filters; 166 | distance_metric?: DistanceMetric; 167 | schema?: Schema; 168 | /** See https://turbopuffer.com/docs/upsert#param-encryption. */ 169 | encryption?: Encryption; 170 | } 171 | 172 | export type QueryRow = RowDoc & 173 | RowAttributes & { 174 | $dist?: number; 175 | }; 176 | 177 | export type QueryResults = { 178 | rows: QueryRow[]; 179 | performance: QueryPerformance; 180 | billing: QueryBilling; 181 | }; 182 | 183 | export interface QueryPerformance extends RequestTiming { 184 | approx_namespace_size: number; 185 | cache_hit_ratio: number; 186 | cache_temperature: string; 187 | server_total_ms: number; 188 | query_execution_ms: number; 189 | exhaustive_search_count: number; 190 | } 191 | 192 | export interface QueryBilling { 193 | billable_logical_bytes_queried: number; 194 | billable_logical_bytes_returned: number; 195 | } 196 | 197 | export interface HintCacheWarmResponse { 198 | message: string; 199 | status: "ACCEPTED" | "OK"; 200 | } 201 | 202 | export interface ExportResponse { 203 | ids: Id[]; 204 | vectors: number[][]; 205 | attributes?: ColumnarAttributes; 206 | next_cursor: string | null; 207 | } 208 | 209 | export interface NamespaceMetadata { 210 | id: string; 211 | approx_count: number; 212 | dimensions: number; 213 | created_at: Date; 214 | } 215 | export interface NamespacesListResult { 216 | namespaces: { id: string }[]; 217 | next_cursor?: string; 218 | } 219 | export interface RecallMeasurement { 220 | avg_recall: number; 221 | avg_exhaustive_count: number; 222 | avg_ann_count: number; 223 | } 224 | -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "rootDir": "src", 4 | "outDir": "dist", 5 | "target": "ES2017", 6 | "module": "commonjs", 7 | "moduleResolution": "Node", 8 | "lib": ["es2017", "es2022", "es7", "es6", "dom.iterable"], 9 | "strict": true, 10 | "sourceMap": true, 11 | "esModuleInterop": true, 12 | "declaration": true, 13 | "declarationMap": true, 14 | "resolveJsonModule": true 15 | }, 16 | "include": ["src"], 17 | "exclude": ["src/*.test.ts"] 18 | } 19 | --------------------------------------------------------------------------------