950a640e
tangwang
embeddings
|
1
|
"""Text embedding client for the local embedding HTTP service."""
|
be52af70
tangwang
first commit
|
2
|
|
950a640e
tangwang
embeddings
|
3
4
5
6
7
|
import logging
import os
import pickle
from datetime import timedelta
from typing import Any, List, Optional, Union
|
be52af70
tangwang
first commit
|
8
|
|
be52af70
tangwang
first commit
|
9
|
import numpy as np
|
453992a8
tangwang
需求:
|
10
|
import redis
|
950a640e
tangwang
embeddings
|
11
|
import requests
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
12
13
|
logger = logging.getLogger(__name__)
|
be52af70
tangwang
first commit
|
14
|
|
42e3aea6
tangwang
tidy
|
15
16
|
from config.services_config import get_embedding_base_url
|
453992a8
tangwang
需求:
|
17
18
19
20
21
22
|
# Try to import REDIS_CONFIG, but allow import to fail
try:
from config.env_config import REDIS_CONFIG
except ImportError:
REDIS_CONFIG = {}
|
be52af70
tangwang
first commit
|
23
|
|
950a640e
tangwang
embeddings
|
24
|
class TextEmbeddingEncoder:
|
be52af70
tangwang
first commit
|
25
|
"""
|
950a640e
tangwang
embeddings
|
26
|
Text embedding encoder using network service.
|
be52af70
tangwang
first commit
|
27
|
"""
|
be52af70
tangwang
first commit
|
28
|
|
950a640e
tangwang
embeddings
|
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
|
def __init__(self, service_url: Optional[str] = None):
resolved_url = service_url or os.getenv("EMBEDDING_SERVICE_URL") or get_embedding_base_url()
self.service_url = str(resolved_url).rstrip("/")
self.endpoint = f"{self.service_url}/embed/text"
self.expire_time = timedelta(days=REDIS_CONFIG.get("cache_expire_days", 180))
logger.info("Creating TextEmbeddingEncoder instance with service URL: %s", self.service_url)
try:
self.redis_client = redis.Redis(
host=REDIS_CONFIG.get("host", "localhost"),
port=REDIS_CONFIG.get("port", 6479),
password=REDIS_CONFIG.get("password"),
decode_responses=False,
socket_timeout=REDIS_CONFIG.get("socket_timeout", 1),
socket_connect_timeout=REDIS_CONFIG.get("socket_connect_timeout", 1),
retry_on_timeout=REDIS_CONFIG.get("retry_on_timeout", False),
health_check_interval=10,
)
self.redis_client.ping()
logger.info("Redis cache initialized for embeddings")
except Exception as e:
logger.warning("Failed to initialize Redis cache for embeddings: %s, continuing without cache", e)
self.redis_client = None
|
be52af70
tangwang
first commit
|
52
|
|
7bfb9946
tangwang
向量化模块
|
53
|
def _call_service(self, request_data: List[str]) -> List[Any]:
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
54
55
56
57
|
"""
Call the embedding service API.
Args:
|
7bfb9946
tangwang
向量化模块
|
58
|
request_data: List of texts
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
59
60
|
Returns:
|
7bfb9946
tangwang
向量化模块
|
61
|
List of embeddings (list[float]) or nulls (None), aligned to input order
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
62
63
64
65
66
67
68
69
70
71
|
"""
try:
response = requests.post(
self.endpoint,
json=request_data,
timeout=60
)
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as e:
|
950a640e
tangwang
embeddings
|
72
|
logger.error(f"TextEmbeddingEncoder service request failed: {e}", exc_info=True)
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
73
74
|
raise
|
be52af70
tangwang
first commit
|
75
76
77
78
|
def encode(
self,
sentences: Union[str, List[str]],
normalize_embeddings: bool = True,
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
79
|
device: str = 'cpu',
|
be52af70
tangwang
first commit
|
80
81
82
|
batch_size: int = 32
) -> np.ndarray:
"""
|
453992a8
tangwang
需求:
|
83
|
Encode text into embeddings via network service with Redis caching.
|
be52af70
tangwang
first commit
|
84
85
86
|
Args:
sentences: Single string or list of strings to encode
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
87
88
89
|
normalize_embeddings: Whether to normalize embeddings (ignored for service)
device: Device parameter ignored for service compatibility
batch_size: Batch size for processing (used for service requests)
|
be52af70
tangwang
first commit
|
90
91
|
Returns:
|
ed948666
tangwang
tidy
|
92
93
|
numpy array of dtype=object,元素均为有效 np.ndarray 向量。
若任一输入无法生成向量,将直接抛出异常。
|
be52af70
tangwang
first commit
|
94
|
"""
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
95
96
97
|
# Convert single string to list
if isinstance(sentences, str):
sentences = [sentences]
|
be52af70
tangwang
first commit
|
98
|
|
453992a8
tangwang
需求:
|
99
|
# Check cache first
|
b2e50710
tangwang
BgeEncoder.encode...
|
100
101
|
uncached_indices: List[int] = []
uncached_texts: List[str] = []
|
453992a8
tangwang
需求:
|
102
|
|
70a318c6
tangwang
fix bug
|
103
104
105
|
embeddings: List[Optional[np.ndarray]] = [None] * len(sentences)
for i, text in enumerate(sentences):
|
950a640e
tangwang
embeddings
|
106
|
cached = self._get_cached_embedding(text, "generic")
|
70a318c6
tangwang
fix bug
|
107
108
109
110
111
112
113
|
if cached is not None:
embeddings[i] = cached
else:
uncached_indices.append(i)
uncached_texts.append(text)
# Prepare request data for uncached texts (after cache check)
|
7bfb9946
tangwang
向量化模块
|
114
|
request_data = list(uncached_texts)
|
453992a8
tangwang
需求:
|
115
116
117
|
# If there are uncached texts, call service
if uncached_texts:
|
ed948666
tangwang
tidy
|
118
|
response_data = self._call_service(request_data)
|
453992a8
tangwang
需求:
|
119
|
|
ed948666
tangwang
tidy
|
120
121
122
123
124
125
126
|
# Process response
for i, text in enumerate(uncached_texts):
original_idx = uncached_indices[i]
if response_data and i < len(response_data):
embedding = response_data[i]
else:
embedding = None
|
7bfb9946
tangwang
向量化模块
|
127
|
|
ed948666
tangwang
tidy
|
128
129
130
131
132
|
if embedding is not None:
embedding_array = np.array(embedding, dtype=np.float32)
if self._is_valid_embedding(embedding_array):
embeddings[original_idx] = embedding_array
self._set_cached_embedding(text, "generic", embedding_array)
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
133
|
else:
|
ed948666
tangwang
tidy
|
134
135
136
137
138
|
raise ValueError(
f"Invalid embedding returned from service for text index {original_idx}"
)
else:
raise ValueError(f"No embedding found for text index {original_idx}: {text[:50]}...")
|
453992a8
tangwang
需求:
|
139
|
|
b2e50710
tangwang
BgeEncoder.encode...
|
140
141
|
# 返回 numpy 数组(dtype=object),元素为 np.ndarray 或 None
return np.array(embeddings, dtype=object)
|
be52af70
tangwang
first commit
|
142
143
144
145
146
|
def encode_batch(
self,
texts: List[str],
batch_size: int = 32,
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
147
|
device: str = 'cpu'
|
be52af70
tangwang
first commit
|
148
149
|
) -> np.ndarray:
"""
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
150
|
Encode a batch of texts efficiently via network service.
|
be52af70
tangwang
first commit
|
151
152
153
154
|
Args:
texts: List of texts to encode
batch_size: Batch size for processing
|
325eec03
tangwang
1. 日志、配置基础设施,使用优化
|
155
|
device: Device parameter ignored for service compatibility
|
be52af70
tangwang
first commit
|
156
157
158
159
160
|
Returns:
numpy array of embeddings
"""
return self.encode(texts, batch_size=batch_size, device=device)
|
453992a8
tangwang
需求:
|
161
162
163
164
165
|
def _get_cache_key(self, query: str, language: str) -> str:
"""Generate a cache key for the query"""
return f"embedding:{language}:{query}"
|
b2e50710
tangwang
BgeEncoder.encode...
|
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
|
def _is_valid_embedding(self, embedding: np.ndarray) -> bool:
"""
Check if embedding is valid (not None, correct shape, no NaN/Inf).
Args:
embedding: Embedding array to validate
Returns:
True if valid, False otherwise
"""
if embedding is None:
return False
if not isinstance(embedding, np.ndarray):
return False
if embedding.size == 0:
return False
# Check for NaN or Inf values
if not np.isfinite(embedding).all():
return False
return True
|
453992a8
tangwang
需求:
|
187
188
189
190
191
192
193
194
195
|
def _get_cached_embedding(self, query: str, language: str) -> Optional[np.ndarray]:
"""Get embedding from cache if exists (with sliding expiration)"""
if not self.redis_client:
return None
try:
cache_key = self._get_cache_key(query, language)
cached_data = self.redis_client.get(cache_key)
if cached_data:
|
b2e50710
tangwang
BgeEncoder.encode...
|
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
|
embedding = pickle.loads(cached_data)
# Validate cached embedding - if invalid, ignore cache and return None
if self._is_valid_embedding(embedding):
logger.debug(f"Cache hit for embedding: {query}")
# Update expiration time on access (sliding expiration)
self.redis_client.expire(cache_key, self.expire_time)
return embedding
else:
logger.warning(
f"Invalid embedding found in cache (contains NaN/Inf or invalid shape), "
f"ignoring cache for query: {query[:50]}..."
)
# Delete invalid cache entry
try:
self.redis_client.delete(cache_key)
except Exception as e:
logger.debug(f"Failed to delete invalid cache entry: {e}")
return None
|
453992a8
tangwang
需求:
|
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
|
return None
except Exception as e:
logger.error(f"Error retrieving embedding from cache: {e}")
return None
def _set_cached_embedding(self, query: str, language: str, embedding: np.ndarray) -> bool:
"""Store embedding in cache"""
if not self.redis_client:
return False
try:
cache_key = self._get_cache_key(query, language)
serialized_data = pickle.dumps(embedding)
self.redis_client.setex(
cache_key,
self.expire_time,
serialized_data
)
logger.debug(f"Successfully cached embedding for query: {query}")
return True
except (redis.exceptions.BusyLoadingError, redis.exceptions.ConnectionError,
redis.exceptions.TimeoutError, redis.exceptions.RedisError) as e:
logger.warning(f"Redis error storing embedding in cache: {e}")
return False
except Exception as e:
logger.error(f"Error storing embedding in cache: {e}")
return False
|