milvus-haystack 0.0.3__tar.gz → 0.0.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,50 @@
1
+ pull_request_rules:
2
+ - name: Add needs-dco label when DCO check failed
3
+ conditions:
4
+ - or:
5
+ - base=main
6
+ - base~=^2(\.\d+){1,2}$
7
+ - -status-success=DCO
8
+ actions:
9
+ label:
10
+ remove:
11
+ - dco-passed
12
+ add:
13
+ - needs-dco
14
+
15
+ - name: Add dco-passed label when DCO check passed
16
+ conditions:
17
+ - or:
18
+ - base=main
19
+ - base~=^2(\.\d+){1,2}$
20
+ - status-success=DCO
21
+ actions:
22
+ label:
23
+ remove:
24
+ - needs-dco
25
+ add:
26
+ - dco-passed
27
+
28
+ - name: Test passed for code changed
29
+ conditions:
30
+ - or:
31
+ - base=main
32
+ - base~=^2(\.\d+){1,2}$
33
+ - 'status-success=Lint and Test / Run Lint and Tests (3.9) (pull_request) '
34
+ actions:
35
+ label:
36
+ add:
37
+ - ci-passed
38
+
39
+ - name: Remove ci-passed label when test failed
40
+ conditions:
41
+ - or:
42
+ - base=main
43
+ - base~=^2(\.\d+){1,2}$
44
+ - or:
45
+ - 'status-success=Lint and Test / Run Lint and Tests (3.9) (pull_request) '
46
+ actions:
47
+ label:
48
+ remove:
49
+ - ci-passed
50
+
@@ -0,0 +1,38 @@
1
+ name: Lint and Test
2
+
3
+ on:
4
+ pull_request:
5
+ paths:
6
+ - "src/**"
7
+ - "tests/**"
8
+ - pyproject.toml
9
+
10
+ concurrency:
11
+ group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
12
+ cancel-in-progress: ${{ github.event_name == 'pull_request' }}
13
+
14
+ jobs:
15
+ build:
16
+ name: Run Lint and Tests
17
+ runs-on: ubuntu-latest
18
+ strategy:
19
+ matrix:
20
+ python-version: ['3.9']
21
+
22
+ steps:
23
+ - name: Checkout code
24
+ uses: actions/checkout@v3
25
+
26
+ - name: Set up Python ${{ matrix.python-version }}
27
+ uses: actions/setup-python@v4
28
+ with:
29
+ python-version: ${{ matrix.python-version }}
30
+
31
+ - name: Install Hatch
32
+ run: pip install --upgrade hatch
33
+
34
+ - name: Lint
35
+ run: hatch run lint:all
36
+
37
+ - name: Run tests
38
+ run: hatch run cov
@@ -0,0 +1,10 @@
1
+ filters:
2
+ ".*":
3
+ reviewers:
4
+ - jaelgu
5
+ - zc277584121
6
+ - codingjaguar
7
+ approvers:
8
+ - jaelgu
9
+ - zc277584121
10
+ - codingjaguar
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: milvus-haystack
3
- Version: 0.0.3
3
+ Version: 0.0.4
4
4
  Project-URL: Documentation, https://github.com/milvus-io/milvus-haystack#readme
5
5
  Project-URL: Issues, https://github.com/milvus-io/milvus-haystack/issues
6
6
  Project-URL: Source, https://github.com/milvus-io/milvus-haystack
@@ -17,7 +17,6 @@ Classifier: Programming Language :: Python :: 3.11
17
17
  Classifier: Programming Language :: Python :: Implementation :: CPython
18
18
  Classifier: Programming Language :: Python :: Implementation :: PyPy
19
19
  Requires-Python: >=3.7
20
- Requires-Dist: farm-haystack
21
20
  Requires-Dist: haystack-ai
22
21
  Requires-Dist: milvus
23
22
  Requires-Dist: pymilvus
@@ -29,7 +29,6 @@ dependencies = [
29
29
  "typing_extensions",
30
30
  "pymilvus",
31
31
  "milvus",
32
- "farm-haystack"
33
32
  ]
34
33
 
35
34
  [project.urls]
@@ -68,7 +67,7 @@ dependencies = [
68
67
  "ruff>=0.0.243",
69
68
  ]
70
69
  [tool.hatch.envs.lint.scripts]
71
- typing = "mypy --install-types --non-interactive {args:src/milvus_haystack tests}"
70
+ typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
72
71
  style = [
73
72
  "ruff {args:.}",
74
73
  "black --check --diff {args:.}",
@@ -175,6 +174,10 @@ markers = [
175
174
  [[tool.mypy.overrides]]
176
175
  module = [
177
176
  "haystack.*",
177
+ "milvus_haystack.*",
178
+ "pymilvus.*",
179
+ "numpy",
180
+ "milvus",
178
181
  "pytest.*"
179
182
  ]
180
183
  ignore_missing_imports = true
@@ -1,4 +1,4 @@
1
1
  # SPDX-FileCopyrightText: 2023-present Tuana Celik <tuana.celik@deepset.ai>
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
- __version__ = "0.0.3"
4
+ __version__ = "0.0.4"
@@ -1,10 +1,7 @@
1
1
  # SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
- from .document_store import MilvusDocumentStore
5
- from .milvus_embedding_retriever import MilvusEmbeddingRetriever
4
+ from .document_store import MilvusDocumentStore # noqa: TID252
5
+ from .milvus_embedding_retriever import MilvusEmbeddingRetriever # noqa: TID252
6
6
 
7
- __all__ = [
8
- "MilvusDocumentStore",
9
- "MilvusEmbeddingRetriever"
10
- ]
7
+ __all__ = ["MilvusDocumentStore", "MilvusEmbeddingRetriever"]
@@ -1,10 +1,12 @@
1
1
  import logging
2
- from typing import List, Dict, Optional, Union, Any
2
+ from typing import Any, Dict, List, Optional, Union
3
3
  from uuid import uuid4
4
+
4
5
  from haystack import Document, default_from_dict, default_to_dict
5
6
  from haystack.document_stores.types import DuplicatePolicy
6
7
  from haystack.errors import FilterError
7
8
  from pymilvus import MilvusException
9
+
8
10
  from milvus_haystack.filters import parse_filters
9
11
 
10
12
  logger = logging.getLogger(__name__)
@@ -24,23 +26,23 @@ class MilvusDocumentStore:
24
26
  """
25
27
 
26
28
  def __init__(
27
- self,
28
- collection_name: str = "HaystackCollection",
29
- collection_description: str = "",
30
- collection_properties: Optional[Dict[str, Any]] = None,
31
- connection_args: Optional[Dict[str, Any]] = None,
32
- consistency_level: str = "Session",
33
- index_params: Optional[dict] = None,
34
- search_params: Optional[dict] = None,
35
- drop_old: Optional[bool] = False,
36
- *,
37
- primary_field: str = "id",
38
- text_field: str = "text",
39
- vector_field: str = "vector",
40
- partition_key_field: Optional[str] = None,
41
- partition_names: Optional[list] = None,
42
- replica_number: int = 1,
43
- timeout: Optional[float] = None,
29
+ self,
30
+ collection_name: str = "HaystackCollection",
31
+ collection_description: str = "",
32
+ collection_properties: Optional[Dict[str, Any]] = None,
33
+ connection_args: Optional[Dict[str, Any]] = None,
34
+ consistency_level: str = "Session",
35
+ index_params: Optional[dict] = None,
36
+ search_params: Optional[dict] = None,
37
+ drop_old: Optional[bool] = False, # noqa: FBT002
38
+ *,
39
+ primary_field: str = "id",
40
+ text_field: str = "text",
41
+ vector_field: str = "vector",
42
+ partition_key_field: Optional[str] = None,
43
+ partition_names: Optional[list] = None,
44
+ replica_number: int = 1,
45
+ timeout: Optional[float] = None,
44
46
  ):
45
47
  """
46
48
  Initialize the Milvus vector store.
@@ -232,9 +234,7 @@ class MilvusDocumentStore:
232
234
  output_fields=output_fields,
233
235
  )
234
236
  except MilvusException as err:
235
- logger.error(
236
- "Failed to query documents with filters expr: %s", expr
237
- )
237
+ logger.error("Failed to query documents with filters expr: %s", expr)
238
238
  raise FilterError(err) from err
239
239
  docs = [self._parse_document(d) for d in res]
240
240
  return docs
@@ -293,7 +293,7 @@ class MilvusDocumentStore:
293
293
  return 0
294
294
 
295
295
  # If the collection hasn't been initialized yet, perform all steps to do so
296
- kwargs = {}
296
+ kwargs: Dict[str, Any] = {}
297
297
  if not isinstance(self.col, Collection):
298
298
  kwargs = {"embeddings": embeddings, "metas": metas}
299
299
  if self.partition_names:
@@ -322,10 +322,9 @@ class MilvusDocumentStore:
322
322
  vectors: list = insert_dict[self._vector_field]
323
323
  total_count = len(vectors)
324
324
 
325
- ids: list[str] = []
326
-
327
325
  batch_size = 1000
328
- assert isinstance(self.col, Collection)
326
+ if not isinstance(self.col, Collection):
327
+ raise MilvusException(message="Collection is not initialized")
329
328
  for i in range(0, total_count, batch_size):
330
329
  # Grab end index
331
330
  end = min(i + batch_size, total_count)
@@ -337,9 +336,7 @@ class MilvusDocumentStore:
337
336
  res = self.col.insert(insert_list, timeout=None, **kwargs)
338
337
  ids.extend(res.primary_keys)
339
338
  except MilvusException as err:
340
- logger.error(
341
- "Failed to insert batch starting at entity: %s/%s", i, total_count
342
- )
339
+ logger.error("Failed to insert batch starting at entity: %s/%s", i, total_count)
343
340
  raise err
344
341
  self.col.flush()
345
342
  return len(ids)
@@ -432,11 +429,11 @@ class MilvusDocumentStore:
432
429
  for con in connections.list_connections():
433
430
  addr = connections.get_connection_addr(con[0])
434
431
  if (
435
- con[1]
436
- and ("address" in addr)
437
- and (addr["address"] == given_address)
438
- and ("user" in addr)
439
- and (addr["user"] == tmp_user)
432
+ con[1]
433
+ and ("address" in addr)
434
+ and (addr["address"] == given_address)
435
+ and ("user" in addr)
436
+ and (addr["user"] == tmp_user)
440
437
  ):
441
438
  logger.debug("Using previous connection: %s", con[0])
442
439
  return con[0]
@@ -452,12 +449,12 @@ class MilvusDocumentStore:
452
449
  raise err
453
450
 
454
451
  def _init(
455
- self,
456
- embeddings: Optional[List] = None,
457
- metas: Optional[List[Dict]] = None,
458
- partition_names: Optional[List] = None,
459
- replica_number: int = 1,
460
- timeout: Optional[float] = None,
452
+ self,
453
+ embeddings: Optional[List] = None,
454
+ metas: Optional[List[Dict]] = None,
455
+ partition_names: Optional[List] = None,
456
+ replica_number: int = 1,
457
+ timeout: Optional[float] = None,
461
458
  ) -> None:
462
459
  if embeddings is not None:
463
460
  self._create_collection(embeddings, metas)
@@ -470,9 +467,7 @@ class MilvusDocumentStore:
470
467
  timeout=timeout,
471
468
  )
472
469
 
473
- def _create_collection(
474
- self, embeddings: list, metas: Optional[List[Dict]] = None
475
- ) -> None:
470
+ def _create_collection(self, embeddings: list, metas: Optional[List[Dict]] = None) -> None:
476
471
  from pymilvus import (
477
472
  Collection,
478
473
  CollectionSchema,
@@ -498,26 +493,16 @@ class MilvusDocumentStore:
498
493
  raise ValueError(err_msg)
499
494
  # Datatype is a string/varchar equivalent
500
495
  elif dtype == DataType.VARCHAR:
501
- fields.append(
502
- FieldSchema(key, DataType.VARCHAR, max_length=65_535)
503
- )
496
+ fields.append(FieldSchema(key, DataType.VARCHAR, max_length=65_535))
504
497
  else:
505
498
  fields.append(FieldSchema(key, dtype))
506
499
 
507
500
  # Create the text field
508
- fields.append(
509
- FieldSchema(self._text_field, DataType.VARCHAR, max_length=65_535)
510
- )
501
+ fields.append(FieldSchema(self._text_field, DataType.VARCHAR, max_length=65_535))
511
502
  # Create the primary key field
512
- fields.append(
513
- FieldSchema(
514
- self._primary_field, DataType.VARCHAR, is_primary=True, max_length=65_535
515
- )
516
- )
503
+ fields.append(FieldSchema(self._primary_field, DataType.VARCHAR, is_primary=True, max_length=65_535))
517
504
  # Create the vector field, supports binary or float vectors
518
- fields.append(
519
- FieldSchema(self._vector_field, infer_dtype_bydata(embeddings[0]), dim=dim)
520
- )
505
+ fields.append(FieldSchema(self._vector_field, infer_dtype_bydata(embeddings[0]), dim=dim))
521
506
 
522
507
  # Create the schema for the collection
523
508
  schema = CollectionSchema(
@@ -538,9 +523,7 @@ class MilvusDocumentStore:
538
523
  if self.collection_properties is not None:
539
524
  self.col.set_properties(self.collection_properties)
540
525
  except MilvusException as err:
541
- logger.error(
542
- "Failed to create collection: %s error: %s", self.collection_name, err
543
- )
526
+ logger.error("Failed to create collection: %s error: %s", self.collection_name, err)
544
527
  raise err
545
528
 
546
529
  def _extract_fields(self) -> None:
@@ -592,9 +575,7 @@ class MilvusDocumentStore:
592
575
  )
593
576
 
594
577
  except MilvusException as err:
595
- logger.error(
596
- "Failed to create an index on collection: %s", self.collection_name
597
- )
578
+ logger.error("Failed to create an index on collection: %s", self.collection_name)
598
579
  raise err
599
580
 
600
581
  def _create_search_params(self) -> None:
@@ -620,20 +601,19 @@ class MilvusDocumentStore:
620
601
  return None
621
602
 
622
603
  def _load(
623
- self,
624
- partition_names: Optional[list] = None,
625
- replica_number: int = 1,
626
- timeout: Optional[float] = None,
604
+ self,
605
+ partition_names: Optional[list] = None,
606
+ replica_number: int = 1,
607
+ timeout: Optional[float] = None,
627
608
  ) -> None:
628
609
  """Load the collection if available."""
629
610
  from pymilvus import Collection, utility
630
611
  from pymilvus.client.types import LoadState
631
612
 
632
613
  if (
633
- isinstance(self.col, Collection)
634
- and self._get_index() is not None
635
- and utility.load_state(self.collection_name, using=self.alias)
636
- == LoadState.NotLoad
614
+ isinstance(self.col, Collection)
615
+ and self._get_index() is not None
616
+ and utility.load_state(self.collection_name, using=self.alias) == LoadState.NotLoad
637
617
  ):
638
618
  self.col.load(
639
619
  partition_names=partition_names,
@@ -642,11 +622,8 @@ class MilvusDocumentStore:
642
622
  )
643
623
 
644
624
  def _embedding_retrieval(
645
- self,
646
- query_embedding: List[float],
647
- filters: Optional[Dict[str, Any]] = None,
648
- top_k: int = 10
649
- ):
625
+ self, query_embedding: List[float], filters: Optional[Dict[str, Any]] = None, top_k: int = 10
626
+ ) -> List[Document]:
650
627
  if self.col is None:
651
628
  logger.debug("No existing collection to search.")
652
629
  return []
@@ -1,4 +1,5 @@
1
- from typing import Dict, Union, Any
1
+ from typing import Any, Dict, Union
2
+
2
3
  from haystack.errors import FilterError
3
4
 
4
5
  LOGIC_OPERATORS = [
@@ -55,10 +56,12 @@ def _parse_comparison(filters: Dict[str, Any]) -> str:
55
56
 
56
57
 
57
58
  def _assert_comparison_filter(filters: Dict[str, Any]):
58
- assert "operator" in filters, "operator must be specified in filters"
59
- assert "field" in filters, "field must be specified in filters"
60
- assert "value" in filters, "value must be specified in filters"
61
- assert filters["operator"] in COMPARISON_OPERATORS, FilterError("operator must be one of: %s" % LOGIC_OPERATORS)
59
+ assert "operator" in filters, "operator must be specified in filters" # noqa: S101
60
+ assert "field" in filters, "field must be specified in filters" # noqa: S101
61
+ assert "value" in filters, "value must be specified in filters" # noqa: S101
62
+ assert filters["operator"] in COMPARISON_OPERATORS, FilterError( # noqa: S101
63
+ "operator must be one of: %s" % LOGIC_OPERATORS
64
+ )
62
65
 
63
66
 
64
67
  def _parse_logic(filters: Dict[str, Any]) -> str:
@@ -80,7 +83,7 @@ def _parse_logic(filters: Dict[str, Any]) -> str:
80
83
 
81
84
 
82
85
  def _assert_logic_filter(filters: Dict[str, Any]):
83
- assert "operator" in filters, "operator must be specified in filters"
84
- assert "conditions" in filters, "conditions must be specified in filters"
85
- assert filters["operator"] in LOGIC_OPERATORS, "operator must be one of: %s" % LOGIC_OPERATORS
86
- assert isinstance(filters["conditions"], list), "conditions must be a list"
86
+ assert "operator" in filters, "operator must be specified in filters" # noqa: S101
87
+ assert "conditions" in filters, "conditions must be specified in filters" # noqa: S101
88
+ assert filters["operator"] in LOGIC_OPERATORS, "operator must be one of: %s" % LOGIC_OPERATORS # noqa: S101
89
+ assert isinstance(filters["conditions"], list), "conditions must be a list" # noqa: S101
@@ -1,5 +1,7 @@
1
- from typing import Any, Dict, Optional, List
2
- from haystack import component, Document
1
+ from typing import Any, Dict, List, Optional
2
+
3
+ from haystack import Document, component
4
+
3
5
  from milvus_haystack import MilvusDocumentStore
4
6
 
5
7
 
@@ -22,7 +24,7 @@ class MilvusEmbeddingRetriever:
22
24
  self.document_store = document_store
23
25
 
24
26
  @component.output_types(documents=List[Document])
25
- def run(self, query_embedding: List[float]):
27
+ def run(self, query_embedding: List[float]) -> Dict[str, List[Document]]:
26
28
  """
27
29
  Retrieve documents from the `MilvusDocumentStore`, based on their dense embeddings.
28
30
 
@@ -1,27 +1,34 @@
1
+ import logging
2
+ import time
3
+
1
4
  import pytest
2
5
  from haystack import Document
3
6
  from haystack.document_stores.types import DocumentStore
4
- from haystack.testing.document_store import CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsTest
7
+ from haystack.testing.document_store import CountDocumentsTest, DeleteDocumentsTest, WriteDocumentsTest
8
+
5
9
  from src.milvus_haystack import MilvusDocumentStore
6
10
 
11
+ logger = logging.getLogger(__name__)
12
+
7
13
 
8
14
  class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsTest):
9
15
  from milvus import MilvusServer
16
+
10
17
  milvus_server = MilvusServer()
11
18
  milvus_server.set_base_dir("test_milvus_base")
12
19
  milvus_server.listen_port = 19530
13
20
  try:
14
21
  milvus_server.stop()
15
- except:
16
- pass
22
+ except Exception as err:
23
+ logger.debug("Can not stop Milvus server. %s", err)
17
24
  try:
18
25
  milvus_server.cleanup()
19
- except:
20
- pass
26
+ except Exception as err:
27
+ logger.debug("Can not cleanup Milvus. %s", err)
21
28
  try:
22
29
  milvus_server.start()
23
- except:
24
- pass
30
+ except Exception as err:
31
+ logger.debug("Can not start Milvus server. %s", err)
25
32
 
26
33
  @pytest.fixture
27
34
  def document_store(self) -> MilvusDocumentStore:
@@ -42,17 +49,26 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
42
49
  )
43
50
  assert document_store.count_documents() == 3
44
51
 
52
+ def test_delete_documents(self, document_store: DocumentStore):
53
+ """
54
+ Test delete_documents() normal behaviour.
55
+ """
56
+ doc = Document(content="test doc")
57
+ document_store.write_documents([doc])
58
+ assert document_store.count_documents() == 1
59
+
60
+ document_store.delete_documents([doc.id])
61
+ time.sleep(1)
62
+ assert document_store.count_documents() == 0
63
+
45
64
  @pytest.mark.skip(reason="Milvus does not currently check if entity primary keys are duplicates")
46
- def test_write_documents_duplicate_fail(self, document_store: DocumentStore):
47
- ...
65
+ def test_write_documents_duplicate_fail(self, document_store: DocumentStore): ...
48
66
 
49
67
  @pytest.mark.skip(reason="Milvus does not currently check if entity primary keys are duplicates")
50
- def test_write_documents_duplicate_skip(self, document_store: DocumentStore):
51
- ...
68
+ def test_write_documents_duplicate_skip(self, document_store: DocumentStore): ...
52
69
 
53
70
  @pytest.mark.skip(reason="Milvus does not currently check if entity primary keys are duplicates")
54
- def test_write_documents_duplicate_overwrite(self, document_store: DocumentStore):
55
- ...
71
+ def test_write_documents_duplicate_overwrite(self, document_store: DocumentStore): ...
56
72
 
57
73
  def test_to_and_from_dict(self, document_store: MilvusDocumentStore):
58
74
  document_store_dict = document_store.to_dict()
@@ -62,13 +78,7 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
62
78
  "collection_name": "HaystackCollection",
63
79
  "collection_description": "",
64
80
  "collection_properties": None,
65
- "connection_args": {
66
- "host": "localhost",
67
- "port": "19530",
68
- "user": "",
69
- "password": "",
70
- "secure": False
71
- },
81
+ "connection_args": {"host": "localhost", "port": "19530", "user": "", "password": "", "secure": False},
72
82
  "consistency_level": "Session",
73
83
  "index_params": None,
74
84
  "search_params": None,
@@ -79,8 +89,8 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
79
89
  "partition_key_field": None,
80
90
  "partition_names": None,
81
91
  "replica_number": 1,
82
- "timeout": None
83
- }
92
+ "timeout": None,
93
+ },
84
94
  }
85
95
  assert document_store_dict == expected_dict
86
96
  reconstructed_document_store = MilvusDocumentStore.from_dict(document_store_dict)
@@ -1,26 +1,32 @@
1
+ import logging
2
+
1
3
  import pytest
2
4
  from haystack import Document
5
+
3
6
  from src.milvus_haystack import MilvusDocumentStore
4
7
  from src.milvus_haystack.milvus_embedding_retriever import MilvusEmbeddingRetriever
5
8
 
9
+ logger = logging.getLogger(__name__)
10
+
6
11
 
7
12
  class TestMilvusEmbeddingTests:
8
13
  from milvus import MilvusServer
14
+
9
15
  milvus_server = MilvusServer()
10
16
  milvus_server.set_base_dir("test_milvus_base")
11
17
  milvus_server.listen_port = 19530
12
18
  try:
13
19
  milvus_server.stop()
14
- except:
15
- pass
20
+ except Exception as err:
21
+ logger.debug("Can not stop Milvus server. %s", err)
16
22
  try:
17
23
  milvus_server.cleanup()
18
- except:
19
- pass
24
+ except Exception as err:
25
+ logger.debug("Can not cleanup Milvus. %s", err)
20
26
  try:
21
27
  milvus_server.start()
22
- except:
23
- pass
28
+ except Exception as err:
29
+ logger.debug("Can not start Milvus server. %s", err)
24
30
 
25
31
  @pytest.fixture
26
32
  def document_store(self) -> MilvusDocumentStore:
@@ -50,7 +56,9 @@ class TestMilvusEmbeddingTests:
50
56
  )
51
57
  documents.append(doc)
52
58
  document_store.write_documents(documents)
53
- retriever = MilvusEmbeddingRetriever(document_store, )
59
+ retriever = MilvusEmbeddingRetriever(
60
+ document_store,
61
+ )
54
62
  query_embedding = [-10.0] * 128
55
63
  res = retriever.run(query_embedding)
56
- assert res["documents"] == doc
64
+ assert res["documents"] == documents
@@ -1,28 +1,34 @@
1
+ import logging
2
+ from typing import List
3
+
1
4
  import numpy as np
2
5
  import pytest
3
- from typing import List
4
6
  from haystack import Document
5
7
  from haystack.testing.document_store import FilterDocumentsTest, _random_embeddings
8
+
6
9
  from src.milvus_haystack import MilvusDocumentStore
7
10
 
11
+ logger = logging.getLogger(__name__)
12
+
8
13
 
9
14
  class TestMilvusFilters(FilterDocumentsTest):
10
15
  from milvus import MilvusServer
16
+
11
17
  milvus_server = MilvusServer()
12
- milvus_server.set_base_dir('test_milvus_base')
18
+ milvus_server.set_base_dir("test_milvus_base")
13
19
  milvus_server.listen_port = 19530
14
20
  try:
15
21
  milvus_server.stop()
16
- except:
17
- pass
22
+ except Exception as err:
23
+ logger.debug("Can not stop Milvus server. %s", err)
18
24
  try:
19
25
  milvus_server.cleanup()
20
- except:
21
- pass
26
+ except Exception as err:
27
+ logger.debug("Can not cleanup Milvus. %s", err)
22
28
  try:
23
29
  milvus_server.start()
24
- except:
25
- pass
30
+ except Exception as err:
31
+ logger.debug("Can not start Milvus server. %s", err)
26
32
 
27
33
  @pytest.fixture
28
34
  def filterable_docs(self) -> List[Document]:
@@ -103,65 +109,49 @@ class TestMilvusFilters(FilterDocumentsTest):
103
109
  assert np.allclose(np.array(r.embedding), np.array(e.embedding), atol=1e-4)
104
110
 
105
111
  @pytest.mark.skip(reason="Milvus doesn't support comparison with dataframe")
106
- def test_comparison_equal_with_dataframe(self, document_store, filterable_docs):
107
- ...
112
+ def test_comparison_equal_with_dataframe(self, document_store, filterable_docs): ...
108
113
 
109
114
  @pytest.mark.skip(reason="Milvus doesn't support comparison with dataframe")
110
- def test_comparison_not_equal_with_dataframe(self, document_store, filterable_docs):
111
- ...
115
+ def test_comparison_not_equal_with_dataframe(self, document_store, filterable_docs): ...
112
116
 
113
117
  @pytest.mark.skip(reason="Milvus doesn't support comparison with dataframe")
114
- def test_comparison_greater_than_with_dataframe(self, document_store, filterable_docs):
115
- ...
118
+ def test_comparison_greater_than_with_dataframe(self, document_store, filterable_docs): ...
116
119
 
117
120
  @pytest.mark.skip(reason="Milvus doesn't support comparison with dataframe")
118
- def test_comparison_greater_than_equal_with_dataframe(self, document_store, filterable_docs):
119
- ...
121
+ def test_comparison_greater_than_equal_with_dataframe(self, document_store, filterable_docs): ...
120
122
 
121
123
  @pytest.mark.skip(reason="Milvus doesn't support comparison with dataframe")
122
- def test_comparison_less_than_with_dataframe(self, document_store, filterable_docs):
123
- ...
124
+ def test_comparison_less_than_with_dataframe(self, document_store, filterable_docs): ...
124
125
 
125
126
  @pytest.mark.skip(reason="Milvus doesn't support comparison with dataframe")
126
- def test_comparison_less_than_equal_with_dataframe(self, document_store, filterable_docs):
127
- ...
127
+ def test_comparison_less_than_equal_with_dataframe(self, document_store, filterable_docs): ...
128
128
 
129
129
  @pytest.mark.skip(reason="Milvus doesn't support comparison with Dates")
130
- def test_comparison_greater_than_with_iso_date(self, document_store, filterable_docs):
131
- ...
130
+ def test_comparison_greater_than_with_iso_date(self, document_store, filterable_docs): ...
132
131
 
133
132
  @pytest.mark.skip(reason="Milvus doesn't support comparison with Dates")
134
- def test_comparison_greater_than_equal_with_iso_date(self, document_store, filterable_docs):
135
- ...
133
+ def test_comparison_greater_than_equal_with_iso_date(self, document_store, filterable_docs): ...
136
134
 
137
135
  @pytest.mark.skip(reason="Milvus doesn't support comparison with Dates")
138
- def test_comparison_less_than_with_iso_date(self, document_store, filterable_docs):
139
- ...
136
+ def test_comparison_less_than_with_iso_date(self, document_store, filterable_docs): ...
140
137
 
141
138
  @pytest.mark.skip(reason="Milvus doesn't support comparison with Dates")
142
- def test_comparison_less_than_equal_with_iso_date(self, document_store, filterable_docs):
143
- ...
139
+ def test_comparison_less_than_equal_with_iso_date(self, document_store, filterable_docs): ...
144
140
 
145
141
  @pytest.mark.skip(reason="Milvus doesn't support comparison with None")
146
- def test_comparison_equal_with_none(self, document_store, filterable_docs):
147
- ...
142
+ def test_comparison_equal_with_none(self, document_store, filterable_docs): ...
148
143
 
149
144
  @pytest.mark.skip(reason="Milvus doesn't support comparison with None")
150
- def test_comparison_not_equal_with_none(self, document_store, filterable_docs):
151
- ...
145
+ def test_comparison_not_equal_with_none(self, document_store, filterable_docs): ...
152
146
 
153
147
  @pytest.mark.skip(reason="Milvus doesn't support comparison with None")
154
- def test_comparison_greater_than_with_none(self, document_store, filterable_docs):
155
- ...
148
+ def test_comparison_greater_than_with_none(self, document_store, filterable_docs): ...
156
149
 
157
150
  @pytest.mark.skip(reason="Milvus doesn't support comparison with None")
158
- def test_comparison_greater_than_equal_with_none(self, document_store, filterable_docs):
159
- ...
151
+ def test_comparison_greater_than_equal_with_none(self, document_store, filterable_docs): ...
160
152
 
161
153
  @pytest.mark.skip(reason="Milvus doesn't support comparison with None")
162
- def test_comparison_less_than_with_none(self, document_store, filterable_docs):
163
- ...
154
+ def test_comparison_less_than_with_none(self, document_store, filterable_docs): ...
164
155
 
165
156
  @pytest.mark.skip(reason="Milvus doesn't support comparison with None")
166
- def test_comparison_less_than_equal_with_none(self, document_store, filterable_docs):
167
- ...
157
+ def test_comparison_less_than_equal_with_none(self, document_store, filterable_docs): ...