Skip to content

Commit 0695581

Browse files
djlidendjliden
authored
Add embeddings-generation and lakebase-pgvector recipes (#49)
Two standalone recipes salvaged from the stalled RAG chat cookbook effort (PR #23). Both stand on their own merits independent of RAG: - embeddings-generation: call a Databricks AI Gateway embedding endpoint via the Databricks SDK. Useful for RAG, semantic search, classification, clustering, and recommendations. - lakebase-pgvector: enable pgvector in Lakebase and run cosine similarity search. Generalized from the RAG-flavored original (schema renamed rag -> vectors, setupRagTables -> setupVectorTables, step numbering fixed) so the recipe applies to any vector search workload. Vector dimension note expanded. Added a prerequisite callout linking to lakebase-create-instance to match the pattern used in lakebase-data-persistence. Registered in recipes array and inserted into recipesInOrder next to the thematically adjacent recipes (embeddings-generation after foundation-models-api; lakebase-pgvector after lakebase-data-persistence).
1 parent 9334a18 commit 0695581

3 files changed

Lines changed: 220 additions & 0 deletions

File tree

‎content/recipes/embeddings-generation.md‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
## Generate Embeddings with AI Gateway
2+
3+
Generate text embeddings from a Databricks AI Gateway endpoint using the Databricks SDK.
4+
5+
### 1. Find an embedding endpoint
6+
7+
```bash
8+
databricks serving-endpoints list --profile <PROFILE>
9+
```
10+
11+
Common embedding endpoints: `databricks-gte-large-en` (1024d), `databricks-bge-large-en` (1024d).
12+
13+
### 2. Configure environment
14+
15+
`.env`:
16+
17+
```bash
18+
DATABRICKS_EMBEDDING_ENDPOINT=databricks-gte-large-en
19+
```
20+
21+
`app.yaml`:
22+
23+
```yaml
24+
env:
25+
- name: DATABRICKS_EMBEDDING_ENDPOINT
26+
value: "databricks-gte-large-en"
27+
```
28+
29+
### 3. Embedding helper
30+
31+
Create `server/lib/embeddings.ts`:
32+
33+
```typescript
34+
import { getWorkspaceClient } from "@databricks/appkit";
35+
36+
const workspaceClient = getWorkspaceClient({});
37+
38+
export async function generateEmbedding(text: string): Promise<number[]> {
39+
const endpoint =
40+
process.env.DATABRICKS_EMBEDDING_ENDPOINT || "databricks-gte-large-en";
41+
const result = await workspaceClient.servingEndpoints.query({
42+
name: endpoint,
43+
input: text,
44+
});
45+
return result.data![0].embedding!;
46+
}
47+
```
48+
49+
No additional dependencies — uses `@databricks/appkit` already in your project.
50+
51+
### 4. Verify
52+
53+
```bash
54+
databricks serving-endpoints query <embedding-endpoint> \
55+
--json '{"input": "Hello, world!"}' \
56+
--profile <PROFILE>
57+
```
58+
59+
Response includes a `data` array with `embedding` (float array).
60+
61+
#### References
62+
63+
- [Query embedding models](https://docs.databricks.com/aws/en/machine-learning/model-serving/query-embedding-models)

‎content/recipes/lakebase-pgvector.md‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
## Lakebase pgvector
2+
3+
Enable vector similarity search in Lakebase using the pgvector extension.
4+
5+
This recipe assumes you have already completed the [Create a Lakebase Instance](/resources/lakebase-create-instance) recipe and have a Lakebase project provisioned.
6+
7+
### 1. Enable pgvector
8+
9+
```bash
10+
databricks psql --project <project-name> --profile <PROFILE> -- -c "
11+
CREATE EXTENSION IF NOT EXISTS vector;
12+
"
13+
```
14+
15+
### 2. Create embedding table
16+
17+
```sql
18+
CREATE SCHEMA IF NOT EXISTS vectors;
19+
20+
CREATE TABLE IF NOT EXISTS vectors.documents (
21+
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
22+
content TEXT NOT NULL,
23+
embedding VECTOR(1024),
24+
metadata JSONB NOT NULL DEFAULT '{}',
25+
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
26+
);
27+
```
28+
29+
> **Vector dimensions**: `VECTOR(1024)` must match your embedding model's output dimension. `databricks-gte-large-en` and `databricks-bge-large-en` both produce 1024 dimensions. If you use a different model (for example, a 768- or 1536-dimension model), change `VECTOR(1024)` to match.
30+
31+
### 3. Server-side vector store module
32+
33+
Create `server/lib/vector-store.ts` with table setup, insert, and similarity search. Call `setupVectorTables(appkit)` from `server.ts` before starting the server.
34+
35+
`server/lib/vector-store.ts`:
36+
37+
```typescript
38+
import type { Application } from "express";
39+
40+
interface AppKitWithLakebase {
41+
lakebase: {
42+
query(
43+
text: string,
44+
params?: unknown[],
45+
): Promise<{ rows: Record<string, unknown>[] }>;
46+
};
47+
server: {
48+
extend(fn: (app: Application) => void): void;
49+
};
50+
}
51+
52+
export async function setupVectorTables(appkit: AppKitWithLakebase) {
53+
try {
54+
await appkit.lakebase.query("CREATE EXTENSION IF NOT EXISTS vector");
55+
} catch (err: unknown) {
56+
const code = (err as { code?: string }).code;
57+
if (code === "42501") {
58+
console.log(
59+
"[vectors] Skipping extension creation — insufficient privileges (likely already exists)",
60+
);
61+
} else {
62+
throw err;
63+
}
64+
}
65+
const { rows } = await appkit.lakebase.query(
66+
`SELECT 1 FROM information_schema.tables
67+
WHERE table_schema = 'vectors' AND table_name = 'documents'`,
68+
);
69+
if (rows.length > 0) return;
70+
await appkit.lakebase.query(`CREATE SCHEMA IF NOT EXISTS vectors`);
71+
await appkit.lakebase.query(`
72+
CREATE TABLE IF NOT EXISTS vectors.documents (
73+
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
74+
content TEXT NOT NULL,
75+
embedding VECTOR(1024),
76+
metadata JSONB NOT NULL DEFAULT '{}',
77+
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
78+
)
79+
`);
80+
}
81+
82+
export async function insertDocument(
83+
appkit: AppKitWithLakebase,
84+
input: {
85+
content: string;
86+
embedding: number[];
87+
metadata?: Record<string, unknown>;
88+
},
89+
) {
90+
const result = await appkit.lakebase.query(
91+
`INSERT INTO vectors.documents (content, embedding, metadata)
92+
VALUES ($1, $2::vector, $3)
93+
RETURNING id, content, metadata, created_at`,
94+
[
95+
input.content,
96+
JSON.stringify(input.embedding),
97+
JSON.stringify(input.metadata ?? {}),
98+
],
99+
);
100+
return result.rows[0];
101+
}
102+
103+
export async function retrieveSimilar(
104+
appkit: AppKitWithLakebase,
105+
queryEmbedding: number[],
106+
limit = 5,
107+
) {
108+
const result = await appkit.lakebase.query(
109+
`SELECT id, content, metadata, 1 - (embedding <=> $1::vector) AS similarity
110+
FROM vectors.documents
111+
WHERE embedding IS NOT NULL
112+
ORDER BY embedding <=> $1::vector
113+
LIMIT $2`,
114+
[JSON.stringify(queryEmbedding), limit],
115+
);
116+
return result.rows;
117+
}
118+
```
119+
120+
> **Distance operators**: `<=>` cosine (default for text), `<->` L2, `<#>` inner product.
121+
122+
### 4. Create an index
123+
124+
Add after inserting initial data (IVFFlat needs representative data to build):
125+
126+
```sql
127+
CREATE INDEX IF NOT EXISTS idx_documents_embedding
128+
ON vectors.documents USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
129+
ANALYZE vectors.documents;
130+
```
131+
132+
> For higher recall without tuning, use `USING hnsw (embedding vector_cosine_ops)` instead.
133+
134+
#### References
135+
136+
- [pgvector](https://github.com/pgvector/pgvector)
137+
- [Lakebase extensions](https://docs.databricks.com/aws/en/oltp/projects/extensions)

‎src/lib/recipes/recipes.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,15 @@ export const recipes: Recipe[] = [
7171
services: ["AI Gateway"],
7272
prerequisites: ["databricks-local-bootstrap"],
7373
},
74+
{
75+
id: "embeddings-generation",
76+
name: "Generate Embeddings with AI Gateway",
77+
description:
78+
"Generate text embeddings from a Databricks AI Gateway endpoint using the Databricks SDK.",
79+
tags: ["AI", "AI Gateway", "Embeddings"],
80+
services: ["AI Gateway"],
81+
prerequisites: ["databricks-local-bootstrap"],
82+
},
7483
{
7584
id: "model-serving-endpoint-creation",
7685
name: "Create a Databricks Model Serving endpoint",
@@ -107,6 +116,15 @@ export const recipes: Recipe[] = [
107116
services: ["Lakebase", "Databricks Apps"],
108117
prerequisites: ["databricks-local-bootstrap", "lakebase-create-instance"],
109118
},
119+
{
120+
id: "lakebase-pgvector",
121+
name: "Lakebase pgvector",
122+
description:
123+
"Enable vector similarity search in Lakebase using the pgvector extension. Covers extension setup, vector table design, insert and cosine retrieval helpers, and IVFFlat/HNSW index options.",
124+
tags: ["Lakebase", "Postgres", "pgvector", "Vector Search", "Embeddings"],
125+
services: ["Lakebase"],
126+
prerequisites: ["databricks-local-bootstrap", "lakebase-create-instance"],
127+
},
110128
{
111129
id: "lakebase-change-data-feed-autoscaling",
112130
name: "Lakebase Change Data Feed: Sync Lakebase to Unity Catalog (Autoscaling)",
@@ -220,7 +238,9 @@ export const recipesInOrder: Recipe[] = [
220238
"databricks-local-bootstrap",
221239
"lakebase-create-instance",
222240
"lakebase-data-persistence",
241+
"lakebase-pgvector",
223242
"foundation-models-api",
243+
"embeddings-generation",
224244
"model-serving-endpoint-creation",
225245
"ai-chat-model-serving",
226246
"lakebase-chat-persistence",

0 commit comments

Comments
 (0)