main: Exposed API for similarity searches with fuzzy + nmslib.

Author: shashank-mugiwara

main.py

@@ -3,11 +3,13 @@
 
 from ingest import ingest_data_router
 from utils import utils_router
+from search import search_router
 
 
 app = FastAPI()
 app.include_router(ingest_data_router)
 app.include_router(utils_router)
+app.include_router(search_router)
 
 if __name__ == '__main__':
     uvicorn.run("main:app", host='0.0.0.0', port=8080, workers=2)

req/KnnSearchIndex.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+from req import KnnSearchVectorFields
+
+
+class KnnSearchIndex(BaseModel):
+    dedupe_vector_nmslib: KnnSearchVectorFields
+

req/KnnSearchQuery.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+from req import KnnSearchIndex
+
+
+class KnnSearchQuery(BaseModel):
+    knn: KnnSearchIndex | None

req/KnnSearchV1.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+from req import KnnSearchQuery
+
+
+class KnnSearchV1(BaseModel):
+    size: int
+    query: KnnSearchQuery
+

req/KnnSearchVectorFields.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class KnnSearchVectorFields(BaseModel):
+    vector: list
+    k: int

req/KnnSimilaritySearch.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+
+class KnnSimilaritySearch(BaseModel):
+    text: str | None
+    size: int
+    k: int

req/__init__.py

@@ -1 +1,6 @@
 from .OpensearchVectorDocumentV1 import OpensearchVectorDocumentV1 as OpensearchVectorDocumentV1
+from .KnnSearchVectorFields import KnnSearchVectorFields
+from .KnnSearchIndex import KnnSearchIndex
+from .KnnSearchQuery import KnnSearchQuery
+from .KnnSearchV1 import KnnSearchV1
+from .KnnSimilaritySearch import KnnSimilaritySearch

search/__init__.py

@@ -0,0 +1,3 @@
+from .search_similarity import get_similar_knn
+from .search_router import router as search_router
+

search/search_router.py

@@ -0,0 +1,54 @@
+import io
+import logging
+from logging.config import dictConfig
+from config import opensearch_config
+from starlette import status
+
+from fastapi.encoders import jsonable_encoder
+from starlette.responses import JSONResponse
+from resp import DocumentInsertResponse
+
+from logger import DedupeKnnLogger
+from req import OpensearchVectorDocumentV1
+from clients import LoadOpenSearchClient
+from fastapi import APIRouter
+import json
+from req import KnnSimilaritySearch, KnnSearchIndex, KnnSearchVectorFields, \
+    KnnSearchV1, KnnSearchQuery
+
+from utils.fast_text_utils import generate_sentence_vector
+from search import get_similar_knn
+
+router = APIRouter()
+opensearch_client = LoadOpenSearchClient().get_opensearch_client()
+config = opensearch_config.OpensearchConfig().get_config()
+
+logger = logging.getLogger("dedupeknn")
+dictConfig(DedupeKnnLogger().dict())
+
+
+@router.post("/api/v1/similarity/knn/search")
+async def get_similar_records_from_opensearch(req: KnnSimilaritySearch):
+    logger.info("Got request for similarity match: {}".format(jsonable_encoder(req)))
+
+    if req.text is None or req.text == "":
+        return JSONResponse(content="Please give valid text.", status_code=status.HTTP_400_BAD_REQUEST)
+    sentence_vector = generate_sentence_vector(req.text)
+    response = await get_similar_knn(vector=sentence_vector.tolist(), k=req.k, size=req.size)
+    response = construct_similarity_response(opensearch_response=response)
+    return JSONResponse(content=response, status_code=status.HTTP_200_OK)
+
+
+def construct_similarity_response(opensearch_response):
+    successful_response = {}
+    hits = []
+    if 'timed_out' in opensearch_response and opensearch_response['timed_out'] == False:
+        if opensearch_response['hits']['total']['value'] > 0:
+            successful_response['total'] = opensearch_response['hits']['total']['value']
+
+            for hit in opensearch_response['hits']['hits']:
+                h = {'id': hit['_id'], 'text': hit['_source']['input_string'], 'score': hit['_score']}
+                hits.append(h)
+
+            successful_response['hits'] = hits
+    return successful_response

search/search_similarity.py

@@ -0,0 +1,22 @@
+from req import KnnSearchV1
+from req import KnnSearchVectorFields
+from req import KnnSearchIndex
+from req import KnnSearchQuery
+from req import KnnSearchV1
+from clients import LoadOpenSearchClient
+from config import OpensearchConfig
+import json
+
+opensearch_client = LoadOpenSearchClient.get_opensearch_client()
+config = OpensearchConfig.get_config()
+
+
+async def get_similar_knn(vector: list, size: int, k: int):
+    knn_fields = KnnSearchVectorFields(vector=vector, k=k)
+    knn_search_index = KnnSearchIndex(dedupe_vector_nmslib=knn_fields)
+    knn_search_query = KnnSearchQuery(knn=knn_search_index)
+    opensearch_search_query = KnnSearchV1(query=knn_search_query, size=size)
+
+    response = await opensearch_client.search(body=json.loads(opensearch_search_query.json()),
+                                              index=config['INDEX_NAME'][0])
+    return response