Blame view

embeddings/text_encoder.py 8.97 KB
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