Add VLLM as vector source

This commit is contained in:
Cohee
2024-06-09 01:03:22 +03:00
parent 1dd21caa66
commit 4e822eeebb
4 changed files with 111 additions and 1 deletions

View File

@@ -16,6 +16,7 @@ const SOURCES = [
'cohere',
'ollama',
'llamacpp',
'vllm',
];
/**
@@ -45,6 +46,8 @@ async function getVector(source, sourceSettings, text, isQuery, directories) {
return require('../vectors/cohere-vectors').getCohereVector(text, isQuery, directories, sourceSettings.model);
case 'llamacpp':
return require('../vectors/llamacpp-vectors').getLlamaCppVector(text, sourceSettings.apiUrl, directories);
case 'vllm':
return require('../vectors/vllm-vectors').getVllmVector(text, sourceSettings.apiUrl, sourceSettings.model, directories);
case 'ollama':
return require('../vectors/ollama-vectors').getOllamaVector(text, sourceSettings.apiUrl, sourceSettings.model, sourceSettings.keep, directories);
}
@@ -91,6 +94,9 @@ async function getBatchVector(source, sourceSettings, texts, isQuery, directorie
case 'llamacpp':
results.push(...await require('../vectors/llamacpp-vectors').getLlamaCppBatchVector(batch, sourceSettings.apiUrl, directories));
break;
case 'vllm':
results.push(...await require('../vectors/vllm-vectors').getVllmBatchVector(batch, sourceSettings.apiUrl, sourceSettings.model, directories));
break;
case 'ollama':
results.push(...await require('../vectors/ollama-vectors').getOllamaBatchVector(batch, sourceSettings.apiUrl, sourceSettings.model, sourceSettings.keep, directories));
break;
@@ -278,6 +284,14 @@ function getSourceSettings(source, request) {
return {
apiUrl: apiUrl,
};
} else if (source === 'vllm') {
const apiUrl = String(request.headers['x-vllm-url']);
const model = String(request.headers['x-vllm-model']);
return {
apiUrl: apiUrl,
model: model,
};
} else if (source === 'ollama') {
const apiUrl = String(request.headers['x-ollama-url']);
const model = String(request.headers['x-ollama-model']);