116 lines
2.7 KiB
JavaScript
116 lines
2.7 KiB
JavaScript
const { QdrantClient } = require('@qdrant/js-client-rest');
|
|
const axios = require('axios');
|
|
require('dotenv').config();
|
|
|
|
const QDRANT_URL = process.env.QDRANT_URL;
|
|
const QDRANT_API_KEY = process.env.QDRANT_API_KEY;
|
|
const OPENAI_API_KEY = process.env.OPENAI_API_KEY;
|
|
const VECTOR_SIZE = 1536; // OpenAI Ada embeddings size
|
|
|
|
const COLLECTION_NAME = 'entities';
|
|
|
|
let client = null;
|
|
|
|
// Initialize Qdrant client
|
|
function initClient() {
|
|
client = new QdrantClient({
|
|
url: QDRANT_URL,
|
|
apiKey: QDRANT_API_KEY
|
|
});
|
|
}
|
|
|
|
// Create or recreate collection
|
|
async function createCollection() {
|
|
if (!client) initClient();
|
|
try {
|
|
await client.recreateCollection(COLLECTION_NAME, {
|
|
vectors: {
|
|
size: VECTOR_SIZE,
|
|
distance: 'Cosine'
|
|
}
|
|
});
|
|
console.log(`Collection '${COLLECTION_NAME}' created/recreated.`);
|
|
} catch (err) {
|
|
console.error('Error creating collection:', err.message);
|
|
throw err;
|
|
}
|
|
}
|
|
|
|
// Get embedding from OpenAI
|
|
async function getEmbedding(text) {
|
|
try {
|
|
const response = await axios.post(
|
|
'https://api.openai.com/v1/embeddings',
|
|
{
|
|
input: text,
|
|
model: 'text-embedding-ada-002'
|
|
},
|
|
{
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
Authorization: `Bearer ${OPENAI_API_KEY}`
|
|
}
|
|
}
|
|
);
|
|
if (response.data && response.data.data && response.data.data[0]) {
|
|
return response.data.data[0].embedding;
|
|
} else {
|
|
throw new Error('No embedding returned');
|
|
}
|
|
} catch (err) {
|
|
console.error('Error getting embedding:', err.message);
|
|
throw err;
|
|
}
|
|
}
|
|
|
|
// Upsert embeddings for an entity
|
|
async function upsertEmbeddings(entity, text) {
|
|
if (!client) initClient();
|
|
try {
|
|
const vector = await getEmbedding(text);
|
|
const point = {
|
|
id: entity,
|
|
vector,
|
|
payload: {
|
|
entity,
|
|
text
|
|
}
|
|
};
|
|
await client.upsertPoints(COLLECTION_NAME, {
|
|
points: [point]
|
|
});
|
|
console.log(`Upserted embeddings for entity '${entity}'.`);
|
|
} catch (err) {
|
|
console.error('Error upserting embeddings:', err.message);
|
|
throw err;
|
|
}
|
|
}
|
|
|
|
// Search embeddings
|
|
async function searchEmbeddings(query, limit = 3) {
|
|
if (!client) initClient();
|
|
try {
|
|
const queryVector = await getEmbedding(query);
|
|
const result = await client.search(COLLECTION_NAME, {
|
|
vector: queryVector,
|
|
limit,
|
|
with_payload: true,
|
|
with_vector: false
|
|
});
|
|
return result.hits.map((hit) => ({
|
|
id: hit.id,
|
|
score: hit.score,
|
|
payload: hit.payload
|
|
}));
|
|
} catch (err) {
|
|
console.error('Error searching embeddings:', err.message);
|
|
throw err;
|
|
}
|
|
}
|
|
|
|
module.exports = {
|
|
initClient,
|
|
createCollection,
|
|
upsertEmbeddings,
|
|
searchEmbeddings
|
|
}; |