diff --git a/semantic_router/index/local.py b/semantic_router/index/local.py index 8f674c62..76d44d82 100644 --- a/semantic_router/index/local.py +++ b/semantic_router/index/local.py @@ -52,7 +52,10 @@ def _remove_and_sync(self, routes_to_delete: dict) -> np.ndarray: for route, utterances in routes_to_delete.items(): # TODO JB: we should be able to vectorize this? for utterance in utterances: - mask &= ~((route_utterances[:, 0] == route) & (route_utterances[:, 1] == utterance)) + mask &= ~( + (route_utterances[:, 0] == route) + & (route_utterances[:, 1] == utterance) + ) # apply the mask to index, routes, and utterances self.index = self.index[mask] self.routes = self.routes[mask] diff --git a/semantic_router/index/postgres.py b/semantic_router/index/postgres.py index 4cdd94f3..6d445f2b 100644 --- a/semantic_router/index/postgres.py +++ b/semantic_router/index/postgres.py @@ -140,8 +140,10 @@ def __init__( try: import psycopg2 except ImportError: - raise ImportError("Please install psycopg2 to use PostgresIndex. " - "You can install it with: `pip install 'semantic-router[postgres]'`") + raise ImportError( + "Please install psycopg2 to use PostgresIndex. " + "You can install it with: `pip install 'semantic-router[postgres]'`" + ) if connection_string: self.connection_string = connection_string else: diff --git a/tests/unit/test_sync.py b/tests/unit/test_sync.py index 3a640f02..2f19236b 100644 --- a/tests/unit/test_sync.py +++ b/tests/unit/test_sync.py @@ -103,6 +103,7 @@ def layer_yaml(): - how are things going? """ + # not all indexes support metadata, so we map the feature here INCLUDE_METADATA_MAP = { PineconeIndex: True, @@ -111,9 +112,12 @@ def layer_yaml(): QdrantIndex: False, PostgresIndex: False, } + + def include_metadata(index_cls): return INCLUDE_METADATA_MAP.get(index_cls, False) + MERGE_FORCE_LOCAL_RESULT_WITH_METADATA = [ Utterance(route="Route 1", utterance="Hello"), Utterance(route="Route 1", utterance="Hi"), @@ -132,6 +136,7 @@ def include_metadata(index_cls): Utterance(route="Route 2", utterance="Hi"), ] + @pytest.fixture def base_encoder(): return DenseEncoder(name="test-encoder", score_threshold=0.5) @@ -367,9 +372,7 @@ def test_auto_sync_merge_force_local( local_utterances = route_layer.index.get_utterances() # we sort to ensure order is the same # TODO JB: there is a bug here where if we include_metadata=True it fails - local_utterances.sort( - key=lambda x: x.to_str(include_metadata=False) - ) + local_utterances.sort(key=lambda x: x.to_str(include_metadata=False)) assert local_utterances == [ Utterance(route="Route 1", utterance="Hello"), Utterance(route="Route 1", utterance="Hi"),