Merge pull request weaviate#367 from weaviate/hybrid_fusion

Add hybrid fusion type

Author: dirkkul

integration/test_graphql.py

@@ -9,6 +9,7 @@
 import weaviate
 from weaviate import Tenant
 from weaviate.data.replication import ConsistencyLevel
+from weaviate.gql.get import HybridFusion
 
 schema = {
     "classes": [
@@ -197,9 +198,14 @@ def test_bm25_no_result(client):
 
 
 @pytest.mark.parametrize("query", ["sponges", "sponges\n"])
-def test_hybrid(client, query: str):
+@pytest.mark.parametrize("fusion_type", [HybridFusion.RANKED, HybridFusion.RELATIVE_SCORE, None])
+def test_hybrid(client, query: str, fusion_type: Optional[HybridFusion]):
     """Test hybrid search with alpha=0.5 to have a combination of BM25 and vector search."""
-    result = client.query.get("Ship", ["name", "description"]).with_hybrid(query, alpha=0.5).do()
+    result = (
+        client.query.get("Ship", ["name", "description"])
+        .with_hybrid(query, alpha=0.5, fusion_type=fusion_type)
+        .do()
+    )
 
     # will find more results. "The Crusty Crab" is still first, because it matches with the BM25 search
     assert len(result["data"]["Get"]["Ship"]) >= 1

test/gql/test_get.py

@@ -6,7 +6,15 @@
 
 from test.util import check_error_message
 from weaviate.data.replication import ConsistencyLevel
-from weaviate.gql.get import GetBuilder, BM25, Hybrid, LinkTo, GroupBy, AdditionalProperties
+from weaviate.gql.get import (
+    GetBuilder,
+    BM25,
+    Hybrid,
+    LinkTo,
+    GroupBy,
+    AdditionalProperties,
+    HybridFusion,
+)
 
 mock_connection_v117 = Mock()
 mock_connection_v117.server_version = "1.17.4"
@@ -70,34 +78,61 @@ def test_get_references(property_name: str, in_class: str, properties: List[str]
 
 
 @pytest.mark.parametrize(
-    "query,vector,alpha,properties,expected",
+    "query,vector,alpha,properties,fusion_type,expected",
     [
         (
             "query",
             [1, 2, 3],
             0.5,
             None,
+            None,
             'hybrid:{query: "query", vector: [1, 2, 3], alpha: 0.5}',
         ),
-        ("query", None, None, None, 'hybrid:{query: "query"}'),
-        ("query", None, None, ["prop1"], 'hybrid:{query: "query", properties: ["prop1"]}'),
+        ("query", None, None, None, None, 'hybrid:{query: "query"}'),
+        ("query", None, None, ["prop1"], None, 'hybrid:{query: "query", properties: ["prop1"]}'),
         (
             "query",
             None,
             None,
             ["prop1", "prop2"],
+            None,
             'hybrid:{query: "query", properties: ["prop1","prop2"]}',
         ),
+        (
+            "query",
+            None,
+            None,
+            None,
+            HybridFusion.RANKED,
+            'hybrid:{query: "query", fusionType: rankedFusion}',
+        ),
+        (
+            "query",
+            None,
+            None,
+            None,
+            HybridFusion.RELATIVE_SCORE,
+            'hybrid:{query: "query", fusionType: relativeScoreFusion}',
+        ),
+        (
+            "query",
+            None,
+            None,
+            None,
+            "relativeScoreFusion",
+            'hybrid:{query: "query", fusionType: relativeScoreFusion}',
+        ),
     ],
 )
 def test_hybrid(
     query: str,
     vector: Optional[List[float]],
     alpha: Optional[float],
     properties: Optional[List[str]],
+    fusion_type: HybridFusion,
     expected: str,
 ):
-    hybrid = Hybrid(query, alpha, vector, properties)
+    hybrid = Hybrid(query, alpha, vector, properties, fusion_type)
     assert str(hybrid) == expected
 
 

weaviate/gql/get.py

@@ -2,6 +2,7 @@
 GraphQL `Get` command.
 """
 from dataclasses import dataclass, Field, fields
+from enum import Enum
 from json import dumps
 from typing import List, Union, Optional, Dict, Tuple
 
@@ -21,7 +22,7 @@
     Sort,
 )
 from weaviate.types import UUID
-from weaviate.util import image_encoder_b64, _capitalize_first_letter, get_valid_uuid
+from weaviate.util import image_encoder_b64, _capitalize_first_letter, get_valid_uuid, BaseEnum
 from weaviate.warnings import _Warnings
 
 try:
@@ -44,12 +45,18 @@ def __str__(self) -> str:
         return "bm25:{" + ret + "}"
 
 
+class HybridFusion(str, BaseEnum):
+    RANKED = "rankedFusion"
+    RELATIVE_SCORE = "relativeScoreFusion"
+
+
 @dataclass
 class Hybrid:
     query: str
     alpha: Optional[float]
     vector: Optional[List[float]]
     properties: Optional[List[str]]
+    fusion_type: Optional[HybridFusion]
 
     def __str__(self) -> str:
         ret = f'query: "{util.strip_newlines(self.query)}"'
@@ -60,6 +67,11 @@ def __str__(self) -> str:
         if self.properties is not None and len(self.properties) > 0:
             props = '","'.join(self.properties)
             ret += f', properties: ["{props}"]'
+        if self.fusion_type is not None:
+            if isinstance(self.fusion_type, Enum):
+                ret += f", fusionType: {self.fusion_type.value}"
+            else:
+                ret += f", fusionType: {self.fusion_type}"
 
         return "hybrid:{" + ret + "}"
 
@@ -1035,6 +1047,7 @@ def with_hybrid(
         alpha: Optional[float] = None,
         vector: Optional[List[float]] = None,
         properties: Optional[List[str]] = None,
+        fusion_type: Optional[HybridFusion] = None,
     ):
         """Get objects using bm25 and vector, then combine the results using a reciprocal ranking algorithm.
 
@@ -1053,8 +1066,10 @@ def with_hybrid(
         properties: Optional[List[str]]:
             Which properties should be searched by BM25. Does not have any effect for vector search. If None or empty
             all properties are searched.
+        fusion_type: Optional[HybridFusionType]:
+            Which fusion type should be used to merge keyword and vector search.
         """
-        self._hybrid = Hybrid(query, alpha, vector, properties)
+        self._hybrid = Hybrid(query, alpha, vector, properties, fusion_type)
         self._contains_filter = True
         return self