Skip to content

Commit f814d13

Browse files
committed
Add Qdrant vector database client
1 parent 09b704a commit f814d13

3 files changed

Lines changed: 317 additions & 0 deletions

File tree

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,7 @@
11
export { Supabase } from "./lib/supabase/supabase.js";
2+
export { Qdrant } from "./lib/qdrant/qdrant.js";
3+
export type {
4+
QdrantDistance,
5+
QdrantPoint,
6+
QdrantPointId,
7+
} from "./lib/qdrant/qdrant.js";
Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
import retry from "retry";
2+
import { config } from "dotenv";
3+
config();
4+
5+
export type QdrantPointId = string | number;
6+
export type QdrantDistance = "Cosine" | "Euclid" | "Dot" | "Manhattan";
7+
8+
export interface QdrantPoint {
9+
id: QdrantPointId;
10+
vector: number[] | Record<string, number[]>;
11+
payload?: Record<string, any>;
12+
}
13+
14+
export interface QdrantRequestOptions {
15+
method?: "GET" | "POST" | "PUT" | "DELETE";
16+
body?: Record<string, any>;
17+
}
18+
19+
export class Qdrant {
20+
QDRANT_URL: string;
21+
QDRANT_API_KEY?: string;
22+
23+
constructor(QDRANT_URL?: string, QDRANT_API_KEY?: string) {
24+
this.QDRANT_URL = (QDRANT_URL || process.env.QDRANT_URL || "").replace(
25+
/\/+$/,
26+
"",
27+
);
28+
this.QDRANT_API_KEY = QDRANT_API_KEY || process.env.QDRANT_API_KEY;
29+
}
30+
31+
createClient() {
32+
if (!this.QDRANT_URL) {
33+
throw new Error("QDRANT_URL is required to create a Qdrant client");
34+
}
35+
return this;
36+
}
37+
38+
async createCollection({
39+
client = this,
40+
collectionName,
41+
vectorSize,
42+
distance = "Cosine",
43+
}: {
44+
client?: Qdrant;
45+
collectionName: string;
46+
vectorSize: number;
47+
distance?: QdrantDistance;
48+
}) {
49+
return client.request(`/collections/${collectionName}`, {
50+
method: "PUT",
51+
body: {
52+
vectors: {
53+
size: vectorSize,
54+
distance,
55+
},
56+
},
57+
});
58+
}
59+
60+
async insertVectorData({
61+
client = this,
62+
collectionName,
63+
points,
64+
}: {
65+
client?: Qdrant;
66+
collectionName: string;
67+
points: QdrantPoint[];
68+
}) {
69+
return client.request(`/collections/${collectionName}/points?wait=true`, {
70+
method: "PUT",
71+
body: { points },
72+
});
73+
}
74+
75+
async getDataFromQuery({
76+
client = this,
77+
collectionName,
78+
vector,
79+
limit = 10,
80+
filter,
81+
withPayload = true,
82+
withVector = false,
83+
}: {
84+
client?: Qdrant;
85+
collectionName: string;
86+
vector: number[] | Record<string, number[]>;
87+
limit?: number;
88+
filter?: Record<string, any>;
89+
withPayload?: boolean;
90+
withVector?: boolean;
91+
}) {
92+
return client.request(`/collections/${collectionName}/points/search`, {
93+
method: "POST",
94+
body: {
95+
vector,
96+
limit,
97+
filter,
98+
with_payload: withPayload,
99+
with_vector: withVector,
100+
},
101+
});
102+
}
103+
104+
async getDataById({
105+
client = this,
106+
collectionName,
107+
ids,
108+
withPayload = true,
109+
withVector = false,
110+
}: {
111+
client?: Qdrant;
112+
collectionName: string;
113+
ids: QdrantPointId[];
114+
withPayload?: boolean;
115+
withVector?: boolean;
116+
}) {
117+
return client.request(`/collections/${collectionName}/points`, {
118+
method: "POST",
119+
body: {
120+
ids,
121+
with_payload: withPayload,
122+
with_vector: withVector,
123+
},
124+
});
125+
}
126+
127+
async updateById({
128+
client = this,
129+
collectionName,
130+
id,
131+
updatedContent,
132+
}: {
133+
client?: Qdrant;
134+
collectionName: string;
135+
id: QdrantPointId;
136+
updatedContent: Record<string, any>;
137+
}) {
138+
return client.request(
139+
`/collections/${collectionName}/points/payload?wait=true`,
140+
{
141+
method: "POST",
142+
body: {
143+
points: [id],
144+
payload: updatedContent,
145+
},
146+
},
147+
);
148+
}
149+
150+
async deleteById({
151+
client = this,
152+
collectionName,
153+
id,
154+
}: {
155+
client?: Qdrant;
156+
collectionName: string;
157+
id: QdrantPointId;
158+
}) {
159+
return client.request(
160+
`/collections/${collectionName}/points/delete?wait=true`,
161+
{
162+
method: "POST",
163+
body: {
164+
points: [id],
165+
},
166+
},
167+
);
168+
}
169+
170+
async request(
171+
path: string,
172+
options: QdrantRequestOptions = {},
173+
): Promise<any> {
174+
return new Promise((resolve, reject) => {
175+
const operation = retry.operation({
176+
retries: 5,
177+
factor: 3,
178+
minTimeout: 1 * 1000,
179+
maxTimeout: 60 * 1000,
180+
randomize: true,
181+
});
182+
183+
operation.attempt(async () => {
184+
try {
185+
const response = await fetch(`${this.QDRANT_URL}${path}`, {
186+
method: options.method || "GET",
187+
headers: {
188+
"Content-Type": "application/json",
189+
...(this.QDRANT_API_KEY
190+
? { "api-key": this.QDRANT_API_KEY }
191+
: {}),
192+
},
193+
body: options.body ? JSON.stringify(options.body) : undefined,
194+
});
195+
const responseBody = await response.json().catch(() => undefined);
196+
197+
if (!response.ok) {
198+
if (response.status >= 500 && operation.retry(new Error())) return;
199+
reject(
200+
new Error(
201+
`Qdrant request failed with status ${response.status}: ${JSON.stringify(
202+
responseBody,
203+
)}`,
204+
),
205+
);
206+
return;
207+
}
208+
209+
resolve(responseBody?.result ?? responseBody);
210+
} catch (error: any) {
211+
if (operation.retry(error)) return;
212+
reject(error);
213+
}
214+
});
215+
});
216+
}
217+
}
Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
import { Qdrant } from "../../lib/qdrant/qdrant";
2+
import { beforeEach, describe, expect, it, vi } from "vitest";
3+
4+
const mockFetch = vi.fn();
5+
6+
global.fetch = mockFetch as any;
7+
8+
describe("Qdrant", () => {
9+
beforeEach(() => {
10+
mockFetch.mockReset();
11+
mockFetch.mockResolvedValue({
12+
ok: true,
13+
json: async () => ({ result: { status: "ok" } }),
14+
});
15+
});
16+
17+
it("upserts vector points through the Qdrant REST API", async () => {
18+
const qdrant = new Qdrant("https://qdrant.example", "secret-key");
19+
20+
const result = await qdrant.insertVectorData({
21+
collectionName: "documents",
22+
points: [
23+
{
24+
id: "doc-1",
25+
vector: [0.1, 0.2, 0.3],
26+
payload: { content: "hello" },
27+
},
28+
],
29+
});
30+
31+
expect(result).toEqual({ status: "ok" });
32+
expect(mockFetch).toHaveBeenCalledWith(
33+
"https://qdrant.example/collections/documents/points?wait=true",
34+
{
35+
method: "PUT",
36+
headers: {
37+
"Content-Type": "application/json",
38+
"api-key": "secret-key",
39+
},
40+
body: JSON.stringify({
41+
points: [
42+
{
43+
id: "doc-1",
44+
vector: [0.1, 0.2, 0.3],
45+
payload: { content: "hello" },
46+
},
47+
],
48+
}),
49+
},
50+
);
51+
});
52+
53+
it("searches a collection with vector query options", async () => {
54+
const qdrant = new Qdrant("https://qdrant.example/");
55+
56+
await qdrant.getDataFromQuery({
57+
collectionName: "documents",
58+
vector: [0.1, 0.2, 0.3],
59+
limit: 3,
60+
filter: { must: [{ key: "namespace", match: { value: "docs" } }] },
61+
withVector: true,
62+
});
63+
64+
expect(mockFetch).toHaveBeenCalledWith(
65+
"https://qdrant.example/collections/documents/points/search",
66+
expect.objectContaining({
67+
method: "POST",
68+
body: JSON.stringify({
69+
vector: [0.1, 0.2, 0.3],
70+
limit: 3,
71+
filter: { must: [{ key: "namespace", match: { value: "docs" } }] },
72+
with_payload: true,
73+
with_vector: true,
74+
}),
75+
}),
76+
);
77+
});
78+
79+
it("throws a useful error when Qdrant returns a non-2xx response", async () => {
80+
mockFetch.mockResolvedValue({
81+
ok: false,
82+
status: 404,
83+
json: async () => ({ status: { error: "Not found" } }),
84+
});
85+
const qdrant = new Qdrant("https://qdrant.example");
86+
87+
await expect(
88+
qdrant.getDataById({
89+
collectionName: "documents",
90+
ids: ["missing"],
91+
}),
92+
).rejects.toThrow("Qdrant request failed with status 404");
93+
});
94+
});

0 commit comments

Comments
 (0)