@gentrace/pinecone
Version:
Gentrace Pinecone v1 plugin for Node.JS
202 lines (201 loc) • 8.29 kB
JavaScript
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
import { init, Pipeline } from "@gentrace/core";
import { FetchInterceptor } from "@mswjs/interceptors/lib/interceptors/fetch";
import { rest } from "msw";
import { setupServer } from "msw/node";
import { DEFAULT_VECTOR } from "../fixtures";
import { initPlugin, Pinecone } from "../index";
function sleep(seconds) {
return __awaiter(this, void 0, void 0, function* () {
return new Promise((resolve) => setTimeout(resolve, seconds * 1000));
});
}
describe("test_pinecone_completion_pipeline", () => {
const completionResponse = {
choices: [{ finish_reason: "stop", index: 0, text: "\n" }],
created: 1682109134,
id: "cmpl-77riQulvtyXo30e14QwSxzGATk2a5",
model: "text-davinci-003",
object: "text_completion",
usage: { completion_tokens: 1, prompt_tokens: 3, total_tokens: 4 },
};
const gentracePipelineRunResponse = {
pipelineRunId: "1f2e8493-5fd1-4359-8cd7-867175d6d9aa",
};
const pineconeProjectName = {
project_name: "example-index",
};
const pineconeFetchResponse = {
vector: [
{
id: "3890",
values: DEFAULT_VECTOR,
metadata: {
testing: "value",
},
},
],
namespace: "example-index",
};
let server;
let interceptor = new FetchInterceptor();
beforeAll(() => {
interceptor.apply();
interceptor.on("request", (request) => {
if (request.url.href === "https://api.pinecone.io/indexes/example-index") {
return request.respondWith({
status: 200,
statusText: "OK",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
name: "example-index",
metric: "cosine",
dimension: 2,
status: { ready: true, state: "Ready" },
host: "example-index-c3q1sek.svc.aped-4627-b74a.pinecone.io",
spec: { serverless: { region: "us-east-1", cloud: "aws" } },
deletion_protection: "disabled",
}),
});
}
else if (request.url.href.startsWith("https://example-index-c3q1sek.svc.aped-4627-b74a.pinecone.io/vectors/fetch?ids=3890")) {
return request.respondWith({
status: 200,
statusText: "OK",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(pineconeFetchResponse),
});
}
else if (request.url.href.startsWith("https://gentrace.ai/api/v1/run")) {
return request.respondWith({
status: 200,
statusText: "OK",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(gentracePipelineRunResponse),
});
}
});
server = setupServer(rest.post("https://api.openai.com/v1/completions", (req, res, ctx) => {
return res(ctx.status(200), ctx.set("Content-Type", "application/json"), ctx.json(completionResponse));
}), rest.post("https://gentrace.ai/api/v1/run", (req, res, ctx) => {
return res(ctx.status(200), ctx.set("Content-Type", "application/json"), ctx.json(gentracePipelineRunResponse));
}));
server.listen();
});
afterAll(() => __awaiter(void 0, void 0, void 0, function* () {
yield sleep(1);
interceptor.dispose();
server.close();
}));
it("should work properly when using Pinecone wrapper simple module", () => __awaiter(void 0, void 0, void 0, function* () {
init({
apiKey: "gentrace-api-key",
});
const pinecone = new Pinecone({
apiKey: "fake-api-key",
});
const index = yield pinecone.index("example-index");
const fetchResponse = yield index.fetch(["3890"], {
gentrace: {
userId: "test-user-id",
},
});
expect(fetchResponse.pipelineRunId).not.toBeDefined();
}));
it("should work properly when using fetch Pinecone wrapper simple module", () => __awaiter(void 0, void 0, void 0, function* () {
init({
apiKey: "gentrace-api-key",
});
const pinecone = new Pinecone({
apiKey: "fake-api-key",
});
const index = yield pinecone.index("example-index");
const fetchResponse = yield index.fetch(["3890"], {
pipelineSlug: "my-slug",
});
expect(fetchResponse.pipelineRunId).toBeDefined();
}));
it("should work properly when using Pinecone wrapper advanced module", () => __awaiter(void 0, void 0, void 0, function* () {
init({
apiKey: "gentrace-api-key",
});
const plugin = yield initPlugin({
apiKey: "fake-api-key",
});
const pipeline = new Pipeline({
slug: "pinecone-index-fetch-pipeline",
plugins: {
pinecone: plugin,
},
});
const runner = pipeline.start();
const pinecone = yield runner.pinecone;
const index = yield pinecone.Index("example-index");
const fetchResponse = yield index.fetch(["3890"], {
pipelineSlug: "my-slug",
});
expect(fetchResponse.pipelineRunId).not.toBeDefined();
}));
it("should work properly when using Pinecone wrapper advanced module (no pipeline run ID)", () => __awaiter(void 0, void 0, void 0, function* () {
init({
apiKey: "gentrace-api-key",
});
const plugin = yield initPlugin({
apiKey: "fake-api-key",
});
const pipeline = new Pipeline({
slug: "pinecone-index-fetch-pipeline",
plugins: {
pinecone: plugin,
},
});
const runner = pipeline.start();
const pinecone = yield runner.pinecone;
const index = yield pinecone.index("example-index");
const fetchResponse = yield index.fetch(["3890"], {
pipelineSlug: "my-slug",
});
expect(fetchResponse.pipelineRunId).not.toBeDefined();
const result = yield runner.submit();
expect(result.pipelineRunId).toBeDefined();
}));
it("should work properly when using Pinecone wrapper advanced module (specify a context)", () => __awaiter(void 0, void 0, void 0, function* () {
init({
apiKey: "gentrace-api-key",
});
const plugin = yield initPlugin({
apiKey: "fake-api-key",
});
const pipeline = new Pipeline({
slug: "pinecone-index-fetch-pipeline",
plugins: {
pinecone: plugin,
},
});
const runner = pipeline.start({
userId: "user-id",
});
const pinecone = yield runner.pinecone;
const index = yield pinecone.index("example-index");
const fetchResponse = yield index.fetch(["3890"], {
pipelineSlug: "my-slug",
});
expect(fetchResponse.pipelineRunId).not.toBeDefined();
const result = yield runner.submit();
expect(result.pipelineRunId).toBeDefined();
}));
});