rara-tools 0.0.11__py3-none-any.whl → 0.0.12__py3-none-any.whl

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.dist-info → rara_tools-0.0.12.dist-info}/METADATA

@@ -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.dist-info → rara_tools-0.0.12.dist-info}/RECORD

@@ -1,7 +1,7 @@
 rara_tools/converters.py,sha256=JcS74VzV6jm12l3C6aqMJBY9nuVW_aevQeCe32KmfrE,1576
 rara_tools/decorators.py,sha256=MjOyvZ5nTkwxwx2JLFEGpKKBysvecFw6EN6UDrSvZLU,2187
 rara_tools/digar_schema_converter.py,sha256=k95U2iRlEA3sh772-v6snhHW6fju6qSTMnvWJ6DpzZk,14254
-rara_tools/elastic.py,sha256=LZfHZqeTDjCEb5YX4CLPJEFffRSZAcRq6AtyP49Fo0E,11575
+rara_tools/elastic.py,sha256=MgPHxZ3UbSTIL8_sT9gU5V4PLKJjo3aQ8CGyhXjRz6M,13065
 rara_tools/exceptions.py,sha256=BwNh4qWxau_ylr9RqZoYwd1KnExI6oWWWDno3jkh8q4,474
 rara_tools/s3.py,sha256=uNDu2HzMYHAWh33RcHeyPFK7gdQfQPxsdfohyIKezEY,4467
 rara_tools/task_reporter.py,sha256=WCcZts9dAUokPc4vbrG3-lNAFLnWaMgE3b3iaUB7mr8,3256
@@ -9,8 +9,8 @@ rara_tools/utils.py,sha256=9vSbmuWYU5ydr4lXBKlUKa0xzDccFsaJv4T-XwgUfuY,2578
 rara_tools/constants/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rara_tools/constants/digitizer.py,sha256=gJ3jOMwuZfKcLqgOAxTyB266VYsskLabJiMUiSz3xX4,297
 rara_tools/constants/general.py,sha256=E9Jaw-YxocS_tOZw9QBoxO3e9KK5EMbLoM0R7D4Iflw,171
-rara_tools-0.0.11.dist-info/LICENSE.md,sha256=hkZVnIZll7e_KNEQzeY94Y9tlzVL8iVZBTMBvDykksU,35142
-rara_tools-0.0.11.dist-info/METADATA,sha256=pDcladCQ1A9O9Wh4UDSh0eHwNyqcGY1BOwPxSJKpLFk,3895
-rara_tools-0.0.11.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-rara_tools-0.0.11.dist-info/top_level.txt,sha256=JwfB5b8BAtW5OFKRln2AQ_WElTRyIBM4nO0FKN1cupY,11
-rara_tools-0.0.11.dist-info/RECORD,,
+rara_tools-0.0.12.dist-info/LICENSE.md,sha256=hkZVnIZll7e_KNEQzeY94Y9tlzVL8iVZBTMBvDykksU,35142
+rara_tools-0.0.12.dist-info/METADATA,sha256=Sz4qM5x__70ANb98HsgCjNjwEsKSumaUcbpGSk5Tfcw,3895
+rara_tools-0.0.12.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+rara_tools-0.0.12.dist-info/top_level.txt,sha256=JwfB5b8BAtW5OFKRln2AQ_WElTRyIBM4nO0FKN1cupY,11
+rara_tools-0.0.12.dist-info/RECORD,,