From f066e504826b54fab4d35b342db9af6f2271cf7d Mon Sep 17 00:00:00 2001
From: Alex Yang <himself65@outlook.com>
Date: Mon, 11 Nov 2024 13:14:08 -0800
Subject: [PATCH] feat: vllm support (#1468)
---
.changeset/pink-dolls-destroy.md | 6 +++++
examples/vllm.ts | 16 +++++++++++
packages/llamaindex/package.json | 3 ++-
packages/llamaindex/src/llm/index.ts | 2 +-
packages/providers/vllm/package.json | 38 +++++++++++++++++++++++++++
packages/providers/vllm/src/index.ts | 1 +
packages/providers/vllm/src/llm.ts | 25 ++++++++++++++++++
packages/providers/vllm/tsconfig.json | 16 +++++++++++
pnpm-lock.yaml | 13 +++++++++
tsconfig.json | 3 +++
10 files changed, 121 insertions(+), 2 deletions(-)
create mode 100644 .changeset/pink-dolls-destroy.md
create mode 100644 examples/vllm.ts
create mode 100644 packages/providers/vllm/package.json
create mode 100644 packages/providers/vllm/src/index.ts
create mode 100644 packages/providers/vllm/src/llm.ts
create mode 100644 packages/providers/vllm/tsconfig.json
diff --git a/.changeset/pink-dolls-destroy.md b/.changeset/pink-dolls-destroy.md
new file mode 100644
index 000000000..33f1e17d8
--- /dev/null
+++ b/.changeset/pink-dolls-destroy.md
@@ -0,0 +1,6 @@
+---
+"llamaindex": patch
+"@llamaindex/vllm": patch
+---
+
+feat: vllm support
diff --git a/examples/vllm.ts b/examples/vllm.ts
new file mode 100644
index 000000000..f65cd375e
--- /dev/null
+++ b/examples/vllm.ts
@@ -0,0 +1,16 @@
+import { VLLM } from "llamaindex";
+
+const llm = new VLLM({
+ model: "NousResearch/Meta-Llama-3-8B-Instruct",
+});
+
+const response = await llm.chat({
+ messages: [
+ {
+ role: "user",
+ content: "Hello?",
+ },
+ ],
+});
+
+console.log(response.message.content);
diff --git a/packages/llamaindex/package.json b/packages/llamaindex/package.json
index 124edc354..d65a4474a 100644
--- a/packages/llamaindex/package.json
+++ b/packages/llamaindex/package.json
@@ -44,6 +44,7 @@
"@llamaindex/portkey-ai": "workspace:*",
"@llamaindex/readers": "workspace:*",
"@llamaindex/replicate": "workspace:*",
+ "@llamaindex/vllm": "workspace:*",
"@mistralai/mistralai": "^1.0.4",
"@mixedbread-ai/sdk": "^2.2.11",
"@pinecone-database/pinecone": "^3.0.2",
@@ -85,10 +86,10 @@
}
},
"devDependencies": {
+ "@huggingface/transformers": "^3.0.2",
"@swc/cli": "^0.5.0",
"@swc/core": "^1.7.22",
"@vercel/postgres": "^0.10.0",
- "@huggingface/transformers": "^3.0.2",
"concurrently": "^9.1.0",
"glob": "^11.0.0",
"pg": "^8.12.0",
diff --git a/packages/llamaindex/src/llm/index.ts b/packages/llamaindex/src/llm/index.ts
index 5596440a0..83e4b7f08 100644
--- a/packages/llamaindex/src/llm/index.ts
+++ b/packages/llamaindex/src/llm/index.ts
@@ -1,3 +1,4 @@
+export { VLLM, type VLLMParams } from "@llamaindex/vllm";
export {
ALL_AVAILABLE_ANTHROPIC_LEGACY_MODELS,
ALL_AVAILABLE_ANTHROPIC_MODELS,
@@ -6,7 +7,6 @@ export {
} from "./anthropic.js";
export { FireworksLLM } from "./fireworks.js";
export { Gemini, GeminiSession } from "./gemini/base.js";
-
export {
GEMINI_MODEL,
type GoogleGeminiSessionOptions,
diff --git a/packages/providers/vllm/package.json b/packages/providers/vllm/package.json
new file mode 100644
index 000000000..b72a2001d
--- /dev/null
+++ b/packages/providers/vllm/package.json
@@ -0,0 +1,38 @@
+{
+ "name": "@llamaindex/vllm",
+ "description": "vLLM Adapter for LlamaIndex",
+ "version": "0.0.1",
+ "type": "module",
+ "main": "./dist/index.cjs",
+ "module": "./dist/index.js",
+ "exports": {
+ ".": {
+ "require": {
+ "types": "./dist/index.d.cts",
+ "default": "./dist/index.cjs"
+ },
+ "import": {
+ "types": "./dist/index.d.ts",
+ "default": "./dist/index.js"
+ }
+ }
+ },
+ "files": [
+ "dist"
+ ],
+ "repository": {
+ "type": "git",
+ "url": "https://github.com/run-llama/LlamaIndexTS.git",
+ "directory": "packages/providers/vllm"
+ },
+ "scripts": {
+ "build": "bunchee",
+ "dev": "bunchee --watch"
+ },
+ "devDependencies": {
+ "bunchee": "5.6.1"
+ },
+ "dependencies": {
+ "@llamaindex/openai": "workspace:*"
+ }
+}
diff --git a/packages/providers/vllm/src/index.ts b/packages/providers/vllm/src/index.ts
new file mode 100644
index 000000000..2c5bc5792
--- /dev/null
+++ b/packages/providers/vllm/src/index.ts
@@ -0,0 +1 @@
+export { VLLM, type VLLMParams } from "./llm";
diff --git a/packages/providers/vllm/src/llm.ts b/packages/providers/vllm/src/llm.ts
new file mode 100644
index 000000000..28b56afc0
--- /dev/null
+++ b/packages/providers/vllm/src/llm.ts
@@ -0,0 +1,25 @@
+/**
+ * vLLM
+ *
+ * https://docs.vllm.ai/en/latest/index.html
+ *
+ * @module
+ */
+import { OpenAI } from "@llamaindex/openai";
+
+export type VLLMParams = {
+ model: string;
+ baseURL?: string;
+};
+
+export class VLLM extends OpenAI {
+ constructor(params: VLLMParams) {
+ super({
+ additionalSessionOptions: {
+ baseURL: "http://localhost:8000/v1",
+ },
+ model: params.model,
+ apiKey: "token-abc123",
+ });
+ }
+}
diff --git a/packages/providers/vllm/tsconfig.json b/packages/providers/vllm/tsconfig.json
new file mode 100644
index 000000000..7af4eb314
--- /dev/null
+++ b/packages/providers/vllm/tsconfig.json
@@ -0,0 +1,16 @@
+{
+ "extends": "../../../tsconfig.json",
+ "compilerOptions": {
+ "target": "ESNext",
+ "module": "ESNext",
+ "moduleResolution": "bundler",
+ "outDir": "./lib",
+ "tsBuildInfoFile": "./lib/.tsbuildinfo"
+ },
+ "include": ["./src", "package.json"],
+ "references": [
+ {
+ "path": "../openai/tsconfig.json"
+ }
+ ]
+}
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 3d650bd2c..c1e41071e 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1026,6 +1026,9 @@ importers:
'@llamaindex/replicate':
specifier: workspace:*
version: link:../providers/replicate
+ '@llamaindex/vllm':
+ specifier: workspace:*
+ version: link:../providers/vllm
'@mistralai/mistralai':
specifier: ^1.0.4
version: 1.1.0(zod@3.23.8)
@@ -1350,6 +1353,16 @@ importers:
specifier: 5.6.1
version: 5.6.1(typescript@5.6.3)
+ packages/providers/vllm:
+ dependencies:
+ '@llamaindex/openai':
+ specifier: workspace:*
+ version: link:../openai
+ devDependencies:
+ bunchee:
+ specifier: 5.6.1
+ version: 5.6.1(typescript@5.6.3)
+
packages/readers:
dependencies:
'@azure/cosmos':
diff --git a/tsconfig.json b/tsconfig.json
index 0d73a1e45..3be6d6b8d 100644
--- a/tsconfig.json
+++ b/tsconfig.json
@@ -53,6 +53,9 @@
{
"path": "./packages/providers/ollama/tsconfig.json"
},
+ {
+ "path": "./packages/providers/vllm/tsconfig.json"
+ },
{
"path": "./packages/cloud/tsconfig.json"
},
--
GitLab