rara-tools 0.0.11__tar.gz → 0.0.12__tar.gz

{rara_tools-0.0.11/rara_tools.egg-info → rara_tools-0.0.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: rara-tools
-Version: 0.0.11
+Version: 0.0.12
 Summary: Tools to support Kata's work.
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10

rara_tools-0.0.12/VERSION

@@ -0,0 +1 @@
+0.0.12

{rara_tools-0.0.11 → rara_tools-0.0.12}/rara_tools/elastic.py

@@ -263,7 +263,7 @@ class KataElastic:
         response = s.execute()
         return response
     
-    def execute_vector_search(
+    def execute_ann_vector_search(
             self,
             index: str,
             field: str,
@@ -281,6 +281,7 @@ class KataElastic:
         :param: query vector List[float]: Vector to search matches for.
         :param: k int: Number of nearest neighbors to return.
         :param: num_candidates int: Number of candidates considered before selecting k results.
+        :param: n_docs: int: Number of documents to return.
         :param: elastic_ids: List[str]: Elastic ID-s for restricting the search.
         """
 
@@ -308,6 +309,48 @@ class KataElastic:
         # Execute the search
         response = s.execute()
         return response
+    
+        
+    def execute_script_score_vector_search(
+            self,
+            index: str,
+            field: str,
+            query_vector: List[float],
+            n_docs: int = 10,
+            elastic_ids: List[str] = []
+    ) -> Response:
+        """ Execute a vector search.
+        NB! Requires different mapping than ANN!
+        
+        :param: index str: Index to search from.
+        :param: field str: Field containing vectorized data.
+        :param: query vector List[float]: Vector to search matches for.
+        :param: n_docs: int: Number of documents to return.
+        :param: elastic_ids: List[str]: Elastic ID-s for restricting the search.
+        """
+        s = elasticsearch_dsl.Search(using=self.elasticsearch, index=index)
+        
+        if elastic_ids:
+            query = elasticsearch_dsl.Q("terms", _id=elastic_ids)
+        else:
+            query = elasticsearch_dsl.Q("match_all")
+        # Apply script_score query
+        s = s.query(
+            "script_score",
+            query=query,
+            script={
+                "source": f"1.0 + cosineSimilarity(params.query_vector, '{field}')",
+                "params": {
+                    "query_vector": query_vector
+                }
+            }
+        )
+        # Set min_score and limit number of documents
+        s = s.extra(size=n_docs)
+
+        # Execute search
+        response = s.execute()
+        return response
 
 
     def __str__(self) -> str:

{rara_tools-0.0.11 → rara_tools-0.0.12/rara_tools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: rara-tools
-Version: 0.0.11
+Version: 0.0.12
 Summary: Tools to support Kata's work.
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10

{rara_tools-0.0.11 → rara_tools-0.0.12}/tests/test_elastic_vector_and_search_operations.py

@@ -15,7 +15,7 @@ TEST_DOCUMENTS = load_json("./tests/test_data/elastic_vectorized_docs.json")
 TEST_VECTOR_DATA = load_json("./tests/test_data/test_vector_data.json")
 TEST_VECTOR = TEST_VECTOR_DATA.get("vector")
     
-es_url = os.getenv("ELASTIC_TEST_URL", "http://localhost:9200")
+es_url = os.getenv("ELASTIC_TEST_URL", "http://rara-elastic.texta.ee:9200")#http://localhost:9200")
 ELASTIC = KataElastic(es_url)
 
 TEST_KNN_INDEX_NAME = "tools_knn_testing_index"
@@ -99,8 +99,8 @@ def test_fuzzy_search():
     
     
 @pytest.mark.order(6)
-def test_vector_search():
-    """ Tests vector search.
+def test_ann_vector_search():
+    """ Tests ANN vector search.
     """
     # Execut fuzzy search to get ID restrictions
     response = ELASTIC.execute_fuzzy_search(
@@ -113,7 +113,7 @@ def test_vector_search():
     assert total_hits == 3    
     elastic_ids = [hit.meta.id for hit in response]
     
-    response = ELASTIC.execute_vector_search(
+    response = ELASTIC.execute_ann_vector_search(
         index=TEST_ANN_INDEX_NAME,
         field="vector",
         query_vector=TEST_VECTOR,
@@ -126,8 +126,35 @@ def test_vector_search():
     assert len(descriptions) == 1
     assert descriptions[0] == "Eesti maletaja ja maleteoreetik"
 
-
+    
 @pytest.mark.order(7)
+def test_script_score_vector_search():
+    """ Tests ANN vector search.
+    """
+    # Execut fuzzy search to get ID restrictions
+    response = ELASTIC.execute_fuzzy_search(
+        index=TEST_KNN_INDEX_NAME,
+        field="variations",
+        entity="Paul Keres",
+        fuzziness=2
+    )
+    total_hits = response.hits.total.value
+    assert total_hits == 3    
+    elastic_ids = [hit.meta.id for hit in response]
+    
+    response = ELASTIC.execute_script_score_vector_search(
+        index=TEST_KNN_INDEX_NAME,
+        field="vector",
+        query_vector=TEST_VECTOR,
+        n_docs=1,
+        elastic_ids=elastic_ids
+    )
+    descriptions = [hit.description for hit in response]
+    assert len(descriptions) == 1
+    assert descriptions[0] == "Eesti maletaja ja maleteoreetik"
+
+    
+@pytest.mark.order(8)
 def test_index_deleting():
     """
     Tests deleting index. We delete the test index now.

rara_tools-0.0.11/VERSION

@@ -1 +0,0 @@
-0.0.11