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.
- milvus_haystack-0.0.4/.github/mergify.yml +50 -0
- milvus_haystack-0.0.4/.github/workflows/test.yml +38 -0
- milvus_haystack-0.0.4/OWNERS +10 -0
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/PKG-INFO +1 -2
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/pyproject.toml +5 -2
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/src/milvus_haystack/__about__.py +1 -1
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/src/milvus_haystack/__init__.py +3 -6
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/src/milvus_haystack/document_store.py +52 -75
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/src/milvus_haystack/filters.py +12 -9
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/src/milvus_haystack/milvus_embedding_retriever.py +5 -3
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/tests/test_document_store.py +32 -22
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/tests/test_embedding_retriever.py +16 -8
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/tests/test_filters.py +30 -40
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/.github/workflows/main.yml +0 -0
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/.gitignore +0 -0
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/LICENSE.txt +0 -0
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/README.md +0 -0
- {milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/tests/__init__.py +0 -0
|
@@ -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
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: milvus-haystack
|
|
3
|
-
Version: 0.0.
|
|
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/
|
|
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,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
|
|
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
|
-
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
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
|
-
|
|
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
|
-
|
|
436
|
-
|
|
437
|
-
|
|
438
|
-
|
|
439
|
-
|
|
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
|
-
|
|
456
|
-
|
|
457
|
-
|
|
458
|
-
|
|
459
|
-
|
|
460
|
-
|
|
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
|
-
|
|
624
|
-
|
|
625
|
-
|
|
626
|
-
|
|
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
|
-
|
|
634
|
-
|
|
635
|
-
|
|
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
|
-
|
|
646
|
-
|
|
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
|
|
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(
|
|
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
|
{milvus_haystack-0.0.3 → milvus_haystack-0.0.4}/src/milvus_haystack/milvus_embedding_retriever.py
RENAMED
|
@@ -1,5 +1,7 @@
|
|
|
1
|
-
from typing import Any, Dict,
|
|
2
|
-
|
|
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,
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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(
|
|
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"] ==
|
|
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(
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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): ...
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|