sie-haystack 0.8.2__tar.gz → 0.9.0__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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: sie-haystack
3
- Version: 0.8.2
3
+ Version: 0.9.0
4
4
  Summary: SIE integration for Haystack
5
5
  Author-email: Superlinked <dev@superlinked.com>
6
6
  License: Apache-2.0
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "sie-haystack"
3
- version = "0.8.2"
3
+ version = "0.9.0"
4
4
  description = "SIE integration for Haystack"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.10"
@@ -6,13 +6,23 @@ and detected objects from text.
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
+ from collections.abc import Mapping
9
10
  from dataclasses import dataclass
10
11
  from typing import Any
11
12
 
12
13
  from haystack import component
14
+ from sie_sdk import RequestError
13
15
  from sie_sdk.types import Item
14
16
 
15
17
 
18
+ def _raise_for_item_error(result: Any) -> None:
19
+ """Raise when SIE reports that extraction failed for this item."""
20
+ error = result.get("error") if isinstance(result, Mapping) else None
21
+ if isinstance(error, Mapping):
22
+ msg = f"Extraction failed: {error.get('message')}"
23
+ raise RequestError(msg, code=error.get("code"), request=result.get("request"))
24
+
25
+
16
26
  @dataclass
17
27
  class Entity:
18
28
  """Extracted entity with position and label information."""
@@ -162,6 +172,7 @@ class SIEExtractor:
162
172
  Item(text=text, metadata=metadata),
163
173
  labels=effective_labels,
164
174
  )
175
+ _raise_for_item_error(result)
165
176
 
166
177
  return {
167
178
  "entities": self._build_entities(result),
@@ -14,10 +14,11 @@ from haystack import Document, component
14
14
  def _scores_by_index(results: Mapping[str, Any], count: int) -> list[float]:
15
15
  """Map ScoreResult entries back to input positions by item_id.
16
16
 
17
- Parses each entry's ``item_id`` defensively: a missing, non-integer,
18
- negative, or out-of-range id is skipped (that document keeps its 0.0
19
- default) so a malformed entry can neither crash the rerank nor mis-assign a
20
- score to the wrong document.
17
+ Each input is sent with ``id=str(position)``, which the server echoes as
18
+ the entry's ``item_id``. Only an exact echo of a sent id is used: an entry
19
+ whose ``item_id`` is missing, not a string, or not a sent id is skipped
20
+ (that document keeps its 0.0 default), so a malformed entry can neither
21
+ crash the rerank nor mis-assign a score to the wrong document.
21
22
 
22
23
  Args:
23
24
  results: ScoreResult envelope from ``SIEClient.score()``.
@@ -26,25 +27,12 @@ def _scores_by_index(results: Mapping[str, Any], count: int) -> list[float]:
26
27
  Returns:
27
28
  Scores indexed by input position (0.0 for any unscored/invalid item).
28
29
  """
30
+ positions = {str(index): index for index in range(count)}
29
31
  scores = [0.0] * count
30
32
  for entry in results.get("scores", []):
31
- item_id = entry.get("item_id", entry.get("index"))
32
- # Accept only a genuine integer or an integer string. Reject bool (an int
33
- # subclass, int(True) == 1), float (int(1.5) == 1), and non-integer
34
- # strings, so a malformed id cannot silently overwrite the wrong position.
35
- if isinstance(item_id, bool):
36
- continue
37
- if isinstance(item_id, int):
38
- idx = item_id
39
- elif isinstance(item_id, str):
40
- try:
41
- idx = int(item_id)
42
- except ValueError:
43
- continue
44
- else:
45
- continue
46
- if 0 <= idx < count:
47
- scores[idx] = float(entry.get("score", 0.0))
33
+ item_id = entry.get("item_id")
34
+ if isinstance(item_id, str) and item_id in positions:
35
+ scores[positions[item_id]] = float(entry.get("score", 0.0))
48
36
  return scores
49
37
 
50
38
 
@@ -141,18 +129,19 @@ class SIERanker:
141
129
 
142
130
  # Prepare items
143
131
  query_item = Item(text=query)
144
- doc_items = [Item(text=doc.content or "") for doc in documents]
132
+ doc_items = [Item(text=doc.content or "", id=str(idx)) for idx, doc in enumerate(documents)]
145
133
 
146
134
  # Score documents
147
135
  results = self.client.score(self._model, query_item, doc_items)
148
136
 
149
137
  # ``results`` is a ScoreResult envelope; ranked entries are under
150
- # ``results["scores"]`` (each a ScoreEntry keyed by ``item_id`` = input
151
- # position, plus ``score``). Map scores back to input order by item_id
152
- # rather than zipping positionally — the entries are sorted by relevance,
153
- # not by input order. The envelope also carries ``results["request"]``
154
- # (request id) and ``results["usage"]`` (token usage); those are
155
- # available but intentionally not surfaced through the Haystack contract.
138
+ # ``results["scores"]`` (each a ScoreEntry whose ``item_id`` echoes the
139
+ # input position sent as the item ``id``, plus ``score``). Map scores
140
+ # back to input order by item_id rather than zipping positionally — the
141
+ # entries are sorted by relevance, not by input order. The envelope also
142
+ # carries ``results["request"]`` (request id) and ``results["usage"]``
143
+ # (token usage); those are available but intentionally not surfaced
144
+ # through the Haystack contract.
156
145
  score_by_index = _scores_by_index(results, len(documents))
157
146
 
158
147
  # Build scored documents
@@ -7,11 +7,12 @@ Haystack-specific helpers.
7
7
  from __future__ import annotations
8
8
 
9
9
  from typing import Any
10
- from unittest.mock import MagicMock
10
+ from unittest.mock import NonCallableMagicMock, create_autospec
11
11
 
12
12
  import numpy as np
13
13
  import pytest
14
14
  from haystack import Document
15
+ from sie_sdk import SIEClient
15
16
 
16
17
  # Default test configuration (matches shared conftest)
17
18
  DEFAULT_EMBEDDING_DIM = 384
@@ -84,6 +85,12 @@ def _create_mock_encode_result(
84
85
  return result
85
86
 
86
87
 
88
+ def _server_item_id(item: Any, index: int) -> str:
89
+ """Return the ``item_id`` the SIE server reports: the sent ``id``, else ``item-<index>``."""
90
+ item_id = item.get("id") if isinstance(item, dict) else None
91
+ return item_id if item_id is not None else f"item-{index}"
92
+
93
+
87
94
  def _create_mock_score_result(query: str, items: list[dict], top_k: int | None = None) -> list[dict[str, Any]]:
88
95
  """Create mock score results."""
89
96
  rng = np.random.default_rng(hash(query) % (2**32))
@@ -127,9 +134,9 @@ def _create_mock_extract_result(text: str, labels: list[str]) -> list[dict[str,
127
134
 
128
135
 
129
136
  @pytest.fixture
130
- def mock_sie_client() -> MagicMock:
137
+ def mock_sie_client() -> NonCallableMagicMock:
131
138
  """Create a mocked SIEClient for unit testing."""
132
- client = MagicMock()
139
+ client = create_autospec(SIEClient, instance=True)
133
140
 
134
141
  def mock_encode(_model: str, items: Any, **kwargs: Any) -> list[dict] | dict:
135
142
  # Check output_types to determine what to include
@@ -159,10 +166,7 @@ def mock_sie_client() -> MagicMock:
159
166
  # Mirror the real SDK: a ScoreResult envelope with ranked entries under
160
167
  # "scores", not a bare list.
161
168
  query_text = _get_text(query)
162
- item_dicts = [
163
- {"id": i.get("id", str(idx)) if isinstance(i, dict) else str(idx), "text": _get_text(i)}
164
- for idx, i in enumerate(items)
165
- ]
169
+ item_dicts = [{"id": _server_item_id(i, idx), "text": _get_text(i)} for idx, i in enumerate(items)]
166
170
  return {
167
171
  "model": _model,
168
172
  "scores": _create_mock_score_result(query_text, item_dicts, kwargs.get("top_k")),
@@ -176,9 +180,9 @@ def mock_sie_client() -> MagicMock:
176
180
  return _create_mock_extract_result(_get_text(items), labels)
177
181
  return [_create_mock_extract_result(_get_text(item), labels) for item in items]
178
182
 
179
- client.encode = MagicMock(side_effect=mock_encode)
180
- client.score = MagicMock(side_effect=mock_score)
181
- client.extract = MagicMock(side_effect=mock_extract)
183
+ client.encode.side_effect = mock_encode
184
+ client.score.side_effect = mock_score
185
+ client.extract.side_effect = mock_extract
182
186
  client.base_url = "http://localhost:8080"
183
187
 
184
188
  return client
@@ -2,8 +2,12 @@
2
2
 
3
3
  from __future__ import annotations
4
4
 
5
+ from unittest.mock import NonCallableMagicMock
6
+
7
+ import pytest
5
8
  from sie_haystack import SIEExtractor
6
9
  from sie_haystack.extractors import Entity
10
+ from sie_sdk import RequestError
7
11
 
8
12
 
9
13
  class TestSIEExtractor:
@@ -116,3 +120,21 @@ class TestSIEExtractor:
116
120
  assert entity.score == 0.95
117
121
  assert entity.start == 0
118
122
  assert entity.end == 10
123
+
124
+
125
+ def test_run_raises_on_item_error(mock_sie_client: NonCallableMagicMock, extract_item_error: dict[str, str]) -> None:
126
+ mock_sie_client.extract.side_effect = None
127
+ mock_sie_client.extract.return_value = {
128
+ "entities": [],
129
+ "relations": [],
130
+ "classifications": [],
131
+ "objects": [],
132
+ "error": dict(extract_item_error),
133
+ }
134
+ extractor = SIEExtractor(model="test-extractor")
135
+ extractor._client = mock_sie_client
136
+
137
+ with pytest.raises(RequestError, match="Extraction failed") as excinfo:
138
+ extractor.run(text="text")
139
+
140
+ assert excinfo.value.code == extract_item_error["code"]
@@ -161,10 +161,7 @@ class TestChromaIntegration:
161
161
 
162
162
  # 1. Set up Chroma (ephemeral for test)
163
163
  _ = chromadb.Client() # Ensure ephemeral client
164
- document_store = ChromaDocumentStore(
165
- collection_name="test_sie_haystack",
166
- embedding_function=None, # We provide embeddings via SIE
167
- )
164
+ document_store = ChromaDocumentStore(collection_name="test_sie_haystack")
168
165
 
169
166
  # 2. Create and embed documents
170
167
  doc_embedder = SIEDocumentEmbedder(base_url=sie_url, model="BAAI/bge-m3")
@@ -3,7 +3,7 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  from collections import Counter
6
- from unittest.mock import MagicMock
6
+ from typing import Any
7
7
 
8
8
  from haystack import Document
9
9
  from sie_haystack import SIERanker
@@ -70,7 +70,8 @@ class TestSIERanker:
70
70
  def test_run_maps_scores_by_item_id(self, mock_sie_client: object) -> None:
71
71
  """Top-ranked document is the relevant one, scores mapped by item_id."""
72
72
  documents = [Document(content=f"doc-{i}") for i in range(5)]
73
- mock_sie_client.score = MagicMock(return_value=_RANKED_ENVELOPE)
73
+ mock_sie_client.score.side_effect = None
74
+ mock_sie_client.score.return_value = _RANKED_ENVELOPE
74
75
  ranker = SIERanker(model="test-reranker")
75
76
  ranker._client = mock_sie_client
76
77
 
@@ -83,7 +84,8 @@ class TestSIERanker:
83
84
  def test_run_skips_malformed_item_id(self, mock_sie_client: object) -> None:
84
85
  """Malformed item_ids are skipped (no crash, no misassignment)."""
85
86
  documents = [Document(content=f"doc-{i}") for i in range(3)]
86
- mock_sie_client.score = MagicMock(return_value=_MALFORMED_ENVELOPE)
87
+ mock_sie_client.score.side_effect = None
88
+ mock_sie_client.score.return_value = _MALFORMED_ENVELOPE
87
89
  ranker = SIERanker(model="test-reranker")
88
90
  ranker._client = mock_sie_client
89
91
 
@@ -94,6 +96,21 @@ class TestSIERanker:
94
96
  assert by_content == {"doc-1": 0.8, "doc-0": 0.0, "doc-2": 0.0}
95
97
  assert ranked[0].content == "doc-1"
96
98
 
99
+ def test_run_ranks_through_sie_client(self, score_stub_server: Any) -> None:
100
+ """Scores returned by a real ``SIEClient.score()`` call rank the matching document first."""
101
+ documents = [
102
+ Document(content="The weather today is sunny with clear skies."),
103
+ Document(content="Python is a popular programming language."),
104
+ Document(content="Nearest neighbor search uses distance metrics."),
105
+ Document(content="Vector similarity search finds similar embeddings."),
106
+ ]
107
+ ranker = SIERanker(base_url=score_stub_server.url, model="test-reranker")
108
+
109
+ ranked = ranker.run(query="vector similarity search", documents=documents)["documents"]
110
+
111
+ assert [d.content for d in ranked] == [documents[i].content for i in (3, 2, 0, 1)]
112
+ assert [d.meta["score"] for d in ranked] == [3.0, 1.0, 0.0, 0.0]
113
+
97
114
  def test_run_empty_list(self, mock_sie_client: object) -> None:
98
115
  """Test that run handles empty document list."""
99
116
  ranker = SIERanker(model="test-reranker")
File without changes
File without changes