document_transformer.py
24.2 KB
1
2
3
4
5
6
7
8
9
10
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
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
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
"""
SPU文档转换器 - 公共转换逻辑。
提取全量和增量索引共用的文档转换逻辑,避免代码冗余。
"""
import pandas as pd
import logging
from typing import Dict, Any, Optional, List
from config import ConfigLoader
logger = logging.getLogger(__name__)
# Try to import translator (optional dependency)
try:
from query.translator import Translator
TRANSLATOR_AVAILABLE = True
except ImportError:
TRANSLATOR_AVAILABLE = False
Translator = None
class SPUDocumentTransformer:
"""SPU文档转换器,将SPU、SKU、Option数据转换为ES文档格式。"""
def __init__(
self,
category_id_to_name: Dict[str, str],
searchable_option_dimensions: List[str],
tenant_config: Optional[Dict[str, Any]] = None,
translator: Optional[Any] = None,
translation_prompts: Optional[Dict[str, str]] = None,
encoder: Optional[Any] = None,
enable_title_embedding: bool = True
):
"""
初始化文档转换器。
Args:
category_id_to_name: 分类ID到名称的映射
searchable_option_dimensions: 可搜索的option维度列表
tenant_config: 租户配置(包含主语言和翻译配置)
translator: 翻译器实例(可选,如果提供则启用翻译功能)
translation_prompts: 翻译提示词配置(可选)
encoder: 文本编码器实例(可选,用于生成title_embedding)
enable_title_embedding: 是否启用标题向量化(默认True)
"""
self.category_id_to_name = category_id_to_name
self.searchable_option_dimensions = searchable_option_dimensions
self.tenant_config = tenant_config or {}
self.translator = translator
self.translation_prompts = translation_prompts or {}
self.encoder = encoder
self.enable_title_embedding = enable_title_embedding
def transform_spu_to_doc(
self,
tenant_id: str,
spu_row: pd.Series,
skus: pd.DataFrame,
options: pd.DataFrame
) -> Optional[Dict[str, Any]]:
"""
将单个SPU行和其SKUs转换为ES文档。
Args:
tenant_id: 租户ID
spu_row: SPU行数据
skus: SKU数据DataFrame
options: Option数据DataFrame
Returns:
ES文档字典
"""
doc = {}
# Tenant ID (required)
doc['tenant_id'] = str(tenant_id)
# SPU ID
spu_id = spu_row['id']
doc['spu_id'] = str(spu_id)
# Validate required fields
if pd.isna(spu_row.get('title')) or not str(spu_row['title']).strip():
logger.error(f"SPU {spu_id} has no title, this may cause search issues")
# 获取租户配置
primary_lang = self.tenant_config.get('primary_language', 'zh')
# 文本字段处理(使用translator的内部逻辑自动处理多语言翻译)
self._fill_text_fields(doc, spu_row, primary_lang)
# 标题向量化处理(如果启用)
if self.enable_title_embedding and self.encoder:
self._fill_title_embedding(doc)
# Tags
if pd.notna(spu_row.get('tags')):
tags_str = str(spu_row['tags'])
doc['tags'] = [tag.strip() for tag in tags_str.split(',') if tag.strip()]
# Category相关字段
self._fill_category_fields(doc, spu_row)
# Option名称(从option表获取)
self._fill_option_names(doc, options)
# Image URL
self._fill_image_url(doc, spu_row)
# Sales (fake_sales)
if pd.notna(spu_row.get('fake_sales')):
try:
doc['sales'] = int(spu_row['fake_sales'])
except (ValueError, TypeError):
doc['sales'] = 0
else:
doc['sales'] = 0
# Process SKUs and build specifications
skus_list, prices, compare_prices, sku_prices, sku_weights, sku_weight_units, total_inventory, specifications = \
self._process_skus(skus, options)
doc['skus'] = skus_list
doc['specifications'] = specifications
# 提取option值(根据配置的searchable_option_dimensions)
self._fill_option_values(doc, skus)
# Calculate price ranges
if prices:
doc['min_price'] = float(min(prices))
doc['max_price'] = float(max(prices))
else:
doc['min_price'] = 0.0
doc['max_price'] = 0.0
if compare_prices:
doc['compare_at_price'] = float(max(compare_prices))
else:
doc['compare_at_price'] = None
# SKU扁平化字段
doc['sku_prices'] = sku_prices
doc['sku_weights'] = sku_weights
doc['sku_weight_units'] = list(set(sku_weight_units)) # 去重
doc['total_inventory'] = total_inventory
# Time fields - convert datetime to ISO format string for ES DATE type
if pd.notna(spu_row.get('create_time')):
create_time = spu_row['create_time']
if hasattr(create_time, 'isoformat'):
doc['create_time'] = create_time.isoformat()
else:
doc['create_time'] = str(create_time)
if pd.notna(spu_row.get('update_time')):
update_time = spu_row['update_time']
if hasattr(update_time, 'isoformat'):
doc['update_time'] = update_time.isoformat()
else:
doc['update_time'] = str(update_time)
return doc
def _fill_text_fields(
self,
doc: Dict[str, Any],
spu_row: pd.Series,
primary_lang: str
):
"""
填充文本字段(根据租户配置处理多语言翻译)。
翻译逻辑:
- 根据 tenant_config 中的 translate_to_zh 和 translate_to_en 决定翻译方向
- 如果 translate_to_zh=true,且店铺语言不是zh,则翻译到中文
- 如果 translate_to_en=true,且店铺语言不是en,则翻译到英文
- 如果两个都是false,则不进行翻译,只填充主语言字段
"""
# 从租户配置中读取翻译方向
translate_to_en = bool(self.tenant_config.get('translate_to_en'))
translate_to_zh = bool(self.tenant_config.get('translate_to_zh'))
# Title
if pd.notna(spu_row.get('title')):
title_text = str(spu_row['title'])
# 使用translator的translate_for_indexing方法,自动处理多语言翻译
if self.translator:
# 根据目标语言选择对应的提示词
prompt_zh = self.translation_prompts.get('product_title_zh') or self.translation_prompts.get('default_zh')
prompt_en = self.translation_prompts.get('product_title_en') or self.translation_prompts.get('default_en')
# 调用translate_for_indexing,自动处理翻译逻辑
translations = self.translator.translate_for_indexing(
title_text,
shop_language=primary_lang,
source_lang=primary_lang,
prompt=prompt_zh if primary_lang == 'zh' else prompt_en,
translate_to_en=translate_to_en,
translate_to_zh=translate_to_zh,
)
# 填充翻译结果
doc['title_zh'] = translations.get('zh') or (title_text if primary_lang == 'zh' else None)
doc['title_en'] = translations.get('en') or (title_text if primary_lang == 'en' else None)
else:
# 无翻译器,只填充主语言字段
if primary_lang == 'zh':
doc['title_zh'] = title_text
doc['title_en'] = None
else:
doc['title_zh'] = None
doc['title_en'] = title_text
else:
doc['title_zh'] = None
doc['title_en'] = None
# Brief
if pd.notna(spu_row.get('brief')):
brief_text = str(spu_row['brief'])
if self.translator:
prompt = self.translation_prompts.get('default_zh') or self.translation_prompts.get('default_en')
translations = self.translator.translate_for_indexing(
brief_text,
shop_language=primary_lang,
source_lang=primary_lang,
prompt=prompt,
translate_to_en=translate_to_en,
translate_to_zh=translate_to_zh,
)
doc['brief_zh'] = translations.get('zh') or (brief_text if primary_lang == 'zh' else None)
doc['brief_en'] = translations.get('en') or (brief_text if primary_lang == 'en' else None)
else:
if primary_lang == 'zh':
doc['brief_zh'] = brief_text
doc['brief_en'] = None
else:
doc['brief_zh'] = None
doc['brief_en'] = brief_text
else:
doc['brief_zh'] = None
doc['brief_en'] = None
# Description
if pd.notna(spu_row.get('description')):
desc_text = str(spu_row['description'])
if self.translator:
prompt = self.translation_prompts.get('default_zh') or self.translation_prompts.get('default_en')
translations = self.translator.translate_for_indexing(
desc_text,
shop_language=primary_lang,
source_lang=primary_lang,
prompt=prompt,
translate_to_en=translate_to_en,
translate_to_zh=translate_to_zh,
)
doc['description_zh'] = translations.get('zh') or (desc_text if primary_lang == 'zh' else None)
doc['description_en'] = translations.get('en') or (desc_text if primary_lang == 'en' else None)
else:
if primary_lang == 'zh':
doc['description_zh'] = desc_text
doc['description_en'] = None
else:
doc['description_zh'] = None
doc['description_en'] = desc_text
else:
doc['description_zh'] = None
doc['description_en'] = None
# Vendor
if pd.notna(spu_row.get('vendor')):
vendor_text = str(spu_row['vendor'])
if self.translator:
prompt = self.translation_prompts.get('default_zh') or self.translation_prompts.get('default_en')
translations = self.translator.translate_for_indexing(
vendor_text,
shop_language=primary_lang,
source_lang=primary_lang,
prompt=prompt,
translate_to_en=translate_to_en,
translate_to_zh=translate_to_zh,
)
doc['vendor_zh'] = translations.get('zh') or (vendor_text if primary_lang == 'zh' else None)
doc['vendor_en'] = translations.get('en') or (vendor_text if primary_lang == 'en' else None)
else:
if primary_lang == 'zh':
doc['vendor_zh'] = vendor_text
doc['vendor_en'] = None
else:
doc['vendor_zh'] = None
doc['vendor_en'] = vendor_text
else:
doc['vendor_zh'] = None
doc['vendor_en'] = None
def _fill_category_fields(self, doc: Dict[str, Any], spu_row: pd.Series):
"""填充类目相关字段。"""
if pd.notna(spu_row.get('category_path')):
category_path = str(spu_row['category_path'])
# 解析category_path - 这是逗号分隔的类目ID列表
category_ids = [cid.strip() for cid in category_path.split(',') if cid.strip()]
# 将ID映射为名称
category_names = []
for cid in category_ids:
if cid in self.category_id_to_name:
category_names.append(self.category_id_to_name[cid])
else:
logger.error(f"Category ID {cid} not found in mapping for SPU {spu_row['id']} (title: {spu_row.get('title', 'N/A')}), category_path={category_path}")
category_names.append(cid) # 使用ID作为备选
# 构建类目路径字符串(用于搜索)
if category_names:
category_path_str = '/'.join(category_names)
doc['category_path_zh'] = category_path_str
doc['category_path_en'] = None # 暂时设为空
# 填充分层类目名称
if len(category_names) > 0:
doc['category1_name'] = category_names[0]
if len(category_names) > 1:
doc['category2_name'] = category_names[1]
if len(category_names) > 2:
doc['category3_name'] = category_names[2]
elif pd.notna(spu_row.get('category')):
# 如果category_path为空,使用category字段作为category1_name的备选
category = str(spu_row['category'])
doc['category_name_zh'] = category
doc['category_name_en'] = None
doc['category_name'] = category
# 尝试从category字段解析多级分类
if '/' in category:
path_parts = category.split('/')
if len(path_parts) > 0:
doc['category1_name'] = path_parts[0].strip()
if len(path_parts) > 1:
doc['category2_name'] = path_parts[1].strip()
if len(path_parts) > 2:
doc['category3_name'] = path_parts[2].strip()
else:
# 如果category不包含"/",直接作为category1_name
doc['category1_name'] = category.strip()
if pd.notna(spu_row.get('category')):
# 确保category相关字段都被设置(如果前面没有设置)
category_name = str(spu_row['category'])
if 'category_name_zh' not in doc:
doc['category_name_zh'] = category_name
if 'category_name_en' not in doc:
doc['category_name_en'] = None
if 'category_name' not in doc:
doc['category_name'] = category_name
if pd.notna(spu_row.get('category_id')):
doc['category_id'] = str(int(spu_row['category_id']))
if pd.notna(spu_row.get('category_level')):
doc['category_level'] = int(spu_row['category_level'])
def _fill_option_names(self, doc: Dict[str, Any], options: pd.DataFrame):
"""填充Option名称字段。"""
if not options.empty:
# 按position排序获取option名称
sorted_options = options.sort_values('position')
if len(sorted_options) > 0 and pd.notna(sorted_options.iloc[0].get('name')):
doc['option1_name'] = str(sorted_options.iloc[0]['name'])
if len(sorted_options) > 1 and pd.notna(sorted_options.iloc[1].get('name')):
doc['option2_name'] = str(sorted_options.iloc[1]['name'])
if len(sorted_options) > 2 and pd.notna(sorted_options.iloc[2].get('name')):
doc['option3_name'] = str(sorted_options.iloc[2]['name'])
def _fill_image_url(self, doc: Dict[str, Any], spu_row: pd.Series):
"""填充图片URL字段。"""
if pd.notna(spu_row.get('image_src')):
image_src = str(spu_row['image_src'])
if not image_src.startswith('http'):
image_src = f"//{image_src}" if image_src.startswith('//') else image_src
doc['image_url'] = image_src
def _process_skus(
self,
skus: pd.DataFrame,
options: pd.DataFrame
) -> tuple:
"""处理SKU数据,返回处理结果。"""
skus_list = []
prices = []
compare_prices = []
sku_prices = []
sku_weights = []
sku_weight_units = []
total_inventory = 0
specifications = []
# 构建option名称映射(position -> name)
option_name_map = {}
if not options.empty:
for _, opt_row in options.iterrows():
position = opt_row.get('position')
name = opt_row.get('name')
if pd.notna(position) and pd.notna(name):
option_name_map[int(position)] = str(name)
for _, sku_row in skus.iterrows():
sku_data = self._transform_sku_row(sku_row, option_name_map)
if sku_data:
skus_list.append(sku_data)
# 收集价格信息
if 'price' in sku_data and sku_data['price'] is not None:
try:
price_val = float(sku_data['price'])
prices.append(price_val)
sku_prices.append(price_val)
except (ValueError, TypeError):
pass
if 'compare_at_price' in sku_data and sku_data['compare_at_price'] is not None:
try:
compare_prices.append(float(sku_data['compare_at_price']))
except (ValueError, TypeError):
pass
# 收集重量信息
if 'weight' in sku_data and sku_data['weight'] is not None:
try:
sku_weights.append(int(float(sku_data['weight'])))
except (ValueError, TypeError):
pass
if 'weight_unit' in sku_data and sku_data['weight_unit']:
sku_weight_units.append(str(sku_data['weight_unit']))
# 收集库存信息
if 'stock' in sku_data and sku_data['stock'] is not None:
try:
total_inventory += int(sku_data['stock'])
except (ValueError, TypeError):
pass
# 构建specifications(从SKU的option值和option表的name)
sku_id = str(sku_row['id'])
if pd.notna(sku_row.get('option1')) and 1 in option_name_map:
specifications.append({
'sku_id': sku_id,
'name': option_name_map[1],
'value': str(sku_row['option1'])
})
if pd.notna(sku_row.get('option2')) and 2 in option_name_map:
specifications.append({
'sku_id': sku_id,
'name': option_name_map[2],
'value': str(sku_row['option2'])
})
if pd.notna(sku_row.get('option3')) and 3 in option_name_map:
specifications.append({
'sku_id': sku_id,
'name': option_name_map[3],
'value': str(sku_row['option3'])
})
return skus_list, prices, compare_prices, sku_prices, sku_weights, sku_weight_units, total_inventory, specifications
def _fill_option_values(self, doc: Dict[str, Any], skus: pd.DataFrame):
"""填充option值字段。"""
option1_values = []
option2_values = []
option3_values = []
for _, sku_row in skus.iterrows():
if pd.notna(sku_row.get('option1')):
option1_values.append(str(sku_row['option1']))
if pd.notna(sku_row.get('option2')):
option2_values.append(str(sku_row['option2']))
if pd.notna(sku_row.get('option3')):
option3_values.append(str(sku_row['option3']))
# 去重并根据配置决定是否写入索引
if 'option1' in self.searchable_option_dimensions:
doc['option1_values'] = list(set(option1_values)) if option1_values else []
else:
doc['option1_values'] = []
if 'option2' in self.searchable_option_dimensions:
doc['option2_values'] = list(set(option2_values)) if option2_values else []
else:
doc['option2_values'] = []
if 'option3' in self.searchable_option_dimensions:
doc['option3_values'] = list(set(option3_values)) if option3_values else []
else:
doc['option3_values'] = []
def _transform_sku_row(self, sku_row: pd.Series, option_name_map: Dict[int, str] = None) -> Optional[Dict[str, Any]]:
"""
将SKU行转换为SKU对象。
Args:
sku_row: SKU行数据
option_name_map: position到option名称的映射
Returns:
SKU字典
"""
sku_data = {}
# SKU ID
sku_data['sku_id'] = str(sku_row['id'])
# Price
if pd.notna(sku_row.get('price')):
try:
sku_data['price'] = float(sku_row['price'])
except (ValueError, TypeError):
sku_data['price'] = None
else:
sku_data['price'] = None
# Compare at price
if pd.notna(sku_row.get('compare_at_price')):
try:
sku_data['compare_at_price'] = float(sku_row['compare_at_price'])
except (ValueError, TypeError):
sku_data['compare_at_price'] = None
else:
sku_data['compare_at_price'] = None
# SKU Code
if pd.notna(sku_row.get('sku')):
sku_data['sku_code'] = str(sku_row['sku'])
# Stock
if pd.notna(sku_row.get('inventory_quantity')):
try:
sku_data['stock'] = int(sku_row['inventory_quantity'])
except (ValueError, TypeError):
sku_data['stock'] = 0
else:
sku_data['stock'] = 0
# Weight
if pd.notna(sku_row.get('weight')):
try:
sku_data['weight'] = float(sku_row['weight'])
except (ValueError, TypeError):
sku_data['weight'] = None
else:
sku_data['weight'] = None
# Weight unit
if pd.notna(sku_row.get('weight_unit')):
sku_data['weight_unit'] = str(sku_row['weight_unit'])
# Option values
if pd.notna(sku_row.get('option1')):
sku_data['option1_value'] = str(sku_row['option1'])
if pd.notna(sku_row.get('option2')):
sku_data['option2_value'] = str(sku_row['option2'])
if pd.notna(sku_row.get('option3')):
sku_data['option3_value'] = str(sku_row['option3'])
# Image src
if pd.notna(sku_row.get('image_src')):
sku_data['image_src'] = str(sku_row['image_src'])
return sku_data
def _fill_title_embedding(self, doc: Dict[str, Any]) -> None:
"""
填充标题向量化字段。
使用英文标题(title_en)生成embedding。如果title_en不存在,则使用title_zh。
Args:
doc: ES文档字典
"""
# 优先使用英文标题,如果没有则使用中文标题
title_text = doc.get('title_en') or doc.get('title_zh')
if not title_text or not title_text.strip():
logger.debug(f"No title text available for embedding, SPU: {doc.get('spu_id')}")
return
try:
# 使用BgeEncoder生成embedding
# encode方法返回numpy数组,形状为(n, 1024)
embeddings = self.encoder.encode(title_text)
if embeddings is not None and len(embeddings) > 0:
# 取第一个embedding(因为只传了一个文本)
embedding = embeddings[0]
# 转换为列表格式(ES需要)
doc['title_embedding'] = embedding.tolist()
logger.debug(f"Generated title_embedding for SPU: {doc.get('spu_id')}, title: {title_text[:50]}...")
else:
logger.warning(f"Failed to generate embedding for title: {title_text[:50]}...")
except Exception as e:
logger.error(f"Error generating title_embedding for SPU {doc.get('spu_id')}: {e}", exc_info=True)