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

src/index.js

@@ -1,9 +1,61 @@
-/**
- * Entry point for the comparison library.
- * Exports the compare function as both named and default export.
- */
-const compare = require('./compare');
+const express = require('express');
+const { fetchEntityData } = require('./tavily');
+const {
+  initClient,
+  createCollection,
+  upsertEmbeddings,
+  searchEmbeddings
+} = require('./qdrant');
+require('dotenv').config();
 
-module.exports = compare;
-module.exports.default = compare;
-module.exports.compare = compare;
+const app = express();
+const PORT = process.env.PORT || 3000;
+
+// Middleware to parse JSON
+app.use(express.json());
+
+// Initialize Qdrant client and collection on startup
+(async () => {
+  try {
+    initClient();
+    await createCollection();
+  } catch (err) {
+    console.error('Failed to initialize Qdrant:', err.message);
+    process.exit(1);
+  }
+})();
+
+// Route to fetch data for an entity and store embeddings
+app.get('/fetch/:entity', async (req, res) => {
+  const entity = req.params.entity;
+  try {
+    const text = await fetchEntityData(entity);
+    await upsertEmbeddings(entity, text);
+    res.json({ status: 'success', entity, textLength: text.length });
+  } catch (err) {
+    res.status(500).json({ status: 'error', message: err.message });
+  }
+});
+
+// Route to search embeddings
+app.get('/search', async (req, res) => {
+  const query = req.query.q;
+  if (!query) {
+    return res.status(400).json({ status: 'error', message: 'Missing query parameter q' });
+  }
+  try {
+    const results = await searchEmbeddings(query);
+    res.json({ status: 'success', query, results });
+  } catch (err) {
+    res.status(500).json({ status: 'error', message: err.message });
+  }
+});
+
+// Health check
+app.get('/', (req, res) => {
+  res.send('Tavily-Qdrant Demo Server');
+});
+
+app.listen(PORT, () => {
+  console.log(`Server running on http://localhost:${PORT}`);
+});

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
+};

src/tavily.js

@@ -0,0 +1,41 @@
+const axios = require('axios');
+require('dotenv').config();
+
+const TAVILY_API_KEY = process.env.TAVILY_API_KEY;
+const TAVILY_ENDPOINT = 'https://api.tavily.com/search';
+
+async function fetchEntityData(entity) {
+  try {
+    const response = await axios.post(
+      TAVILY_ENDPOINT,
+      {
+        query: entity,
+        search_depth: 2,
+        include_raw_content: true,
+        max_results: 5
+      },
+      {
+        headers: {
+          'Content-Type': 'application/json',
+          'accept': 'application/json',
+          'Authorization': `Bearer ${TAVILY_API_KEY}`
+        }
+      }
+    );
+
+    if (response.data && response.data.results) {
+      // Concatenate all raw content into a single string
+      const texts = response.data.results
+        .map((r) => r.raw_content || '')
+        .filter(Boolean);
+      return texts.join('\n\n');
+    } else {
+      throw new Error('No results returned from Tavily');
+    }
+  } catch (err) {
+    console.error('Error fetching data from Tavily:', err.message);
+    throw err;
+  }
+}
+
+module.exports = { fetchEntityData };