Ver Fonte

implement appendVectorScore

Yuki Takei há 1 ano atrás
pai
commit
1d16f59998

+ 4 - 4
apps/app/src/server/routes/search.ts

@@ -109,7 +109,7 @@ module.exports = function(crowi, app) {
   api.search = async function(req, res) {
     const user = req.user;
     const {
-      q = null, nq = null, type = null, sort = null, order = null,
+      q = null, nq = null, type = null, sort = null, order = null, vector = null,
     } = req.query;
     let paginateOpts;
 
@@ -135,15 +135,15 @@ module.exports = function(crowi, app) {
     ] : null;
 
     const searchOpts = {
-      ...paginateOpts, type, sort, order,
+      ...paginateOpts, type, sort, order, vector,
     };
 
     let searchResult;
     let delegatorName;
     try {
-      const keyword = decodeURIComponent(q);
+      const query = decodeURIComponent(q);
       const nqName = nq ?? decodeURIComponent(nq);
-      [searchResult, delegatorName] = await searchService.searchKeyword(keyword, nqName, user, userGroups, searchOpts);
+      [searchResult, delegatorName] = await searchService.searchKeyword(query, nqName, user, userGroups, searchOpts);
     }
     catch (err) {
       logger.error('Failed to search', err);

+ 22 - 1
apps/app/src/server/service/search-delegator/elasticsearch.ts

@@ -847,6 +847,20 @@ class ElasticsearchDelegator implements SearchDelegator<Data, ESTermsKey, ESQuer
     };
   }
 
+  async appendVectorScore(query, queryString: string, username: string): Promise<void> {
+    const queryVector = (await openaiService.embed(username, queryString))[0].embedding;
+
+    query.body.query = {
+      script_score: {
+        query: { ...query.body.query },
+        script: {
+          source: "cosineSimilarity(params.query_vector, doc['body_embedded']) + 1.0",
+          params: { query_vector: queryVector },
+        },
+      },
+    };
+  }
+
   appendHighlight(query) {
     query.body.highlight = {
       fragmenter: 'simple',
@@ -877,6 +891,7 @@ class ElasticsearchDelegator implements SearchDelegator<Data, ESTermsKey, ESQuer
     const size = option?.limit ?? null;
     const sort = option?.sort ?? null;
     const order = option?.order ?? null;
+
     const query = this.createSearchQuery();
     this.appendCriteriaForQueryString(query, terms);
 
@@ -885,7 +900,13 @@ class ElasticsearchDelegator implements SearchDelegator<Data, ESTermsKey, ESQuer
     this.appendResultSize(query, from, size);
 
     this.appendSortOrder(query, sort, order);
-    await this.appendFunctionScore(query, queryString);
+
+    if (option?.vector) {
+      await this.appendVectorScore(query, queryString, user.username);
+    }
+    else {
+      await this.appendFunctionScore(query, queryString);
+    }
 
     this.appendHighlight(query);