feat: solution for 'Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily)'

src/qdrant.js

@@ -0,0 +1,116 @@
+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
+};