950a640e
tangwang
embeddings
|
1
2
3
4
|
import pickle
from typing import Any, Dict, List, Optional
import numpy as np
|
ed948666
tangwang
tidy
|
5
|
import pytest
|
950a640e
tangwang
embeddings
|
6
7
8
9
10
|
from config import (
FunctionScoreConfig,
IndexConfig,
QueryConfig,
|
950a640e
tangwang
embeddings
|
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
|
RerankConfig,
SPUConfig,
SearchConfig,
)
from embeddings.text_encoder import TextEmbeddingEncoder
from query import QueryParser
class _FakeRedis:
def __init__(self):
self.store: Dict[str, bytes] = {}
def ping(self):
return True
def get(self, key: str):
return self.store.get(key)
def setex(self, key: str, _expire, value: bytes):
self.store[key] = value
return True
def expire(self, key: str, _expire):
return key in self.store
def delete(self, key: str):
self.store.pop(key, None)
return True
class _FakeResponse:
def __init__(self, payload: List[Optional[List[float]]]):
self._payload = payload
def raise_for_status(self):
return None
def json(self):
return self._payload
class _FakeTranslator:
def translate(
self,
text: str,
target_lang: str,
source_lang: Optional[str] = None,
prompt: Optional[str] = None,
) -> str:
return f"{text}-{target_lang}"
class _FakeQueryEncoder:
def encode(self, sentences, **kwargs):
if isinstance(sentences, str):
sentences = [sentences]
return np.array([np.array([0.11, 0.22, 0.33], dtype=np.float32) for _ in sentences], dtype=object)
def _build_test_config() -> SearchConfig:
return SearchConfig(
field_boosts={"title.en": 3.0},
indexes=[IndexConfig(name="default", label="default", fields=["title.en"], boost=1.0)],
query_config=QueryConfig(
supported_languages=["en", "zh"],
default_language="en",
enable_text_embedding=True,
enable_query_rewrite=False,
|
950a640e
tangwang
embeddings
|
79
|
rewrite_dictionary={},
|
950a640e
tangwang
embeddings
|
80
81
82
|
text_embedding_field="title_embedding",
image_embedding_field=None,
),
|
77ab67ad
tangwang
更新测试用例
|
83
|
function_score=FunctionScoreConfig(),
|
950a640e
tangwang
embeddings
|
84
85
86
87
88
89
90
91
92
93
94
95
96
|
rerank=RerankConfig(),
spu_config=SPUConfig(enabled=True, spu_field="spu_id", inner_hits_size=3),
es_index_name="test_products",
tenant_config={},
es_settings={},
services={},
)
def test_text_embedding_encoder_response_alignment(monkeypatch):
fake_redis = _FakeRedis()
monkeypatch.setattr("embeddings.text_encoder.redis.Redis", lambda **kwargs: fake_redis)
|
77ab67ad
tangwang
更新测试用例
|
97
|
def _fake_post(url, json, timeout, **kwargs):
|
950a640e
tangwang
embeddings
|
98
99
|
assert url.endswith("/embed/text")
assert json == ["hello", "world"]
|
ed948666
tangwang
tidy
|
100
|
return _FakeResponse([[0.1, 0.2], [0.3, 0.4]])
|
950a640e
tangwang
embeddings
|
101
102
103
104
105
106
107
108
109
|
monkeypatch.setattr("embeddings.text_encoder.requests.post", _fake_post)
encoder = TextEmbeddingEncoder(service_url="http://127.0.0.1:6005")
out = encoder.encode(["hello", "world"])
assert len(out) == 2
assert isinstance(out[0], np.ndarray)
assert out[0].shape == (2,)
|
ed948666
tangwang
tidy
|
110
111
112
113
114
115
116
117
|
assert isinstance(out[1], np.ndarray)
assert out[1].shape == (2,)
def test_text_embedding_encoder_raises_on_missing_vector(monkeypatch):
fake_redis = _FakeRedis()
monkeypatch.setattr("embeddings.text_encoder.redis.Redis", lambda **kwargs: fake_redis)
|
77ab67ad
tangwang
更新测试用例
|
118
|
def _fake_post(url, json, timeout, **kwargs):
|
ed948666
tangwang
tidy
|
119
120
121
122
123
124
125
|
return _FakeResponse([[0.1, 0.2], None])
monkeypatch.setattr("embeddings.text_encoder.requests.post", _fake_post)
encoder = TextEmbeddingEncoder(service_url="http://127.0.0.1:6005")
with pytest.raises(ValueError):
encoder.encode(["hello", "world"])
|
950a640e
tangwang
embeddings
|
126
127
128
129
130
131
132
133
134
135
|
def test_text_embedding_encoder_cache_hit(monkeypatch):
fake_redis = _FakeRedis()
cached = np.array([0.9, 0.8], dtype=np.float32)
fake_redis.store["embedding:generic:cached-text"] = pickle.dumps(cached)
monkeypatch.setattr("embeddings.text_encoder.redis.Redis", lambda **kwargs: fake_redis)
calls = {"count": 0}
|
77ab67ad
tangwang
更新测试用例
|
136
|
def _fake_post(url, json, timeout, **kwargs):
|
950a640e
tangwang
embeddings
|
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
|
calls["count"] += 1
return _FakeResponse([[0.3, 0.4]])
monkeypatch.setattr("embeddings.text_encoder.requests.post", _fake_post)
encoder = TextEmbeddingEncoder(service_url="http://127.0.0.1:6005")
out = encoder.encode(["cached-text", "new-text"])
assert calls["count"] == 1
assert np.allclose(out[0], cached)
assert np.allclose(out[1], np.array([0.3, 0.4], dtype=np.float32))
def test_query_parser_generates_query_vector_with_encoder():
parser = QueryParser(
config=_build_test_config(),
text_encoder=_FakeQueryEncoder(),
translator=_FakeTranslator(),
)
parsed = parser.parse("red dress", tenant_id="162", generate_vector=True)
assert parsed.query_vector is not None
assert parsed.query_vector.shape == (3,)
def test_query_parser_skips_query_vector_when_disabled():
parser = QueryParser(
config=_build_test_config(),
text_encoder=_FakeQueryEncoder(),
translator=_FakeTranslator(),
)
parsed = parser.parse("red dress", tenant_id="162", generate_vector=False)
assert parsed.query_vector is None
|