n8n-nodes-databricks
Version:
Databricks node for n8n
114 lines • 4.72 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.DatabricksVectorStoreLangChain = void 0;
const vectorstores_1 = require("@langchain/core/vectorstores");
const documents_1 = require("@langchain/core/documents");
class DatabricksVectorStoreLangChain extends vectorstores_1.VectorStore {
_vectorstoreType() { return "databricks"; }
constructor(embeddings, config) {
super(embeddings, {});
this.config = config;
}
async makeRequest(method, indexName, body) {
const headers = {
'Authorization': `Bearer ${this.config.token}`,
'Content-Type': 'application/json;charset=UTF-8',
'Accept': 'application/json, text/plain, */*',
};
const url = `${this.config.workspaceUrl}/api/2.0/vector-search/indexes/${indexName}/query`;
const response = await fetch(url, {
method,
headers,
body: body ? JSON.stringify(body) : undefined,
});
if (!response.ok) {
const error = await response.json();
throw new Error(`Databricks API error: ${error.message}`);
}
return response.json();
}
static async fromDocuments(docs, embeddings, config) {
const instance = new this(embeddings, config);
await instance.addDocuments(docs);
return instance;
}
static async fromExistingIndex(embeddings, config) {
return new this(embeddings, config);
}
async addDocuments(documents) {
const texts = documents.map((doc) => doc.pageContent);
const vectors = await this.embeddings.embedDocuments(texts);
await this.addVectors(vectors, documents);
}
async addVectors(vectors, documents) {
const rows = vectors.map((vector, i) => ({
id: documents[i].metadata.id || `doc_${i}`,
embedding: vector,
[this.config.textColumn]: documents[i].pageContent,
...Object.fromEntries(this.config.metadataColumns.map(col => [col, documents[i].metadata[col]])),
}));
await this.makeRequest('POST', this.config.indexName, {
vectors: rows,
});
}
async delete(params) {
await this.makeRequest('POST', this.config.indexName, {
ids: params.ids,
});
}
async similaritySearchVectorWithScore(query, k, filterJson, queryType, extraColumns, scoreThreshold) {
let normalizedQuery = query;
if (Array.isArray(query) && query.length === 1 && typeof query[0] === 'object' && query[0] !== null && 'response' in query[0]) {
normalizedQuery = query[0].response;
}
const columns = [this.config.textColumn, ...this.config.metadataColumns];
if (extraColumns) {
for (const col of extraColumns) {
if (!columns.includes(col))
columns.push(col);
}
}
const body = {
columns,
num_results: k,
query_vector: normalizedQuery
};
if (filterJson) {
body.filters_json = filterJson;
}
if (queryType) {
body.query_type = queryType;
}
if (scoreThreshold !== undefined) {
body.score_threshold = scoreThreshold;
}
else if (this.config.scoreThreshold !== undefined) {
body.score_threshold = this.config.scoreThreshold;
}
const response = await this.makeRequest('POST', this.config.indexName, body);
if (!(response === null || response === void 0 ? void 0 : response.result)) {
throw new Error(`Databricks API returned invalid response structure. Full response: ${JSON.stringify(response)}`);
}
if (!response.result.data_array || !Array.isArray(response.result.data_array) || response.result.data_array.length === 0) {
response.result.data_array = [];
}
return response.result.data_array.map(([id, text, vector]) => {
const doc = new documents_1.Document({
pageContent: text,
metadata: {
id,
...(this.config.metadataColumns.length > 0 && {
...Object.fromEntries(this.config.metadataColumns.map((col, index) => [
col,
response.result.data_array[index + 3]
]))
})
},
});
const score = 1.0;
return [doc, score];
});
}
}
exports.DatabricksVectorStoreLangChain = DatabricksVectorStoreLangChain;
//# sourceMappingURL=DatabricksVectorStoreLangChain.js.map