Blame view

search/searcher.py 19.2 KB
be52af70   tangwang   first commit
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
  """
  Main Searcher module - executes search queries against Elasticsearch.
  
  Handles query parsing, boolean expressions, ranking, and result formatting.
  """
  
  from typing import Dict, Any, List, Optional
  import time
  
  from config import CustomerConfig
  from utils.es_client import ESClient
  from query import QueryParser, ParsedQuery
  from indexer import MappingGenerator
  from .boolean_parser import BooleanParser, QueryNode
  from .es_query_builder import ESQueryBuilder
b926f678   tangwang   多语言查询
16
  from .multilang_query_builder import MultiLanguageQueryBuilder
be52af70   tangwang   first commit
17
  from .ranking_engine import RankingEngine
16c42787   tangwang   feat: implement r...
18
  from context.request_context import RequestContext, RequestContextStage, create_request_context
be52af70   tangwang   first commit
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
  
  
  class SearchResult:
      """Container for search results."""
  
      def __init__(
          self,
          hits: List[Dict[str, Any]],
          total: int,
          max_score: float,
          took_ms: int,
          aggregations: Optional[Dict[str, Any]] = None,
          query_info: Optional[Dict[str, Any]] = None
      ):
          self.hits = hits
          self.total = total
          self.max_score = max_score
          self.took_ms = took_ms
          self.aggregations = aggregations or {}
          self.query_info = query_info or {}
  
      def to_dict(self) -> Dict[str, Any]:
          """Convert to dictionary representation."""
          return {
              "hits": self.hits,
              "total": self.total,
              "max_score": self.max_score,
              "took_ms": self.took_ms,
              "aggregations": self.aggregations,
              "query_info": self.query_info
          }
  
  
  class Searcher:
      """
      Main search engine class.
  
      Handles:
      - Query parsing and translation
      - Boolean expression parsing
      - ES query building
      - Result ranking and formatting
      """
  
      def __init__(
          self,
          config: CustomerConfig,
          es_client: ESClient,
          query_parser: Optional[QueryParser] = None
      ):
          """
          Initialize searcher.
  
          Args:
              config: Customer configuration
              es_client: Elasticsearch client
              query_parser: Query parser (created if not provided)
          """
          self.config = config
          self.es_client = es_client
          self.query_parser = query_parser or QueryParser(config)
  
          # Initialize components
          self.boolean_parser = BooleanParser()
          self.ranking_engine = RankingEngine(config.ranking.expression)
  
          # Get mapping info
          mapping_gen = MappingGenerator(config)
          self.match_fields = mapping_gen.get_match_fields_for_domain("default")
          self.text_embedding_field = mapping_gen.get_text_embedding_field()
          self.image_embedding_field = mapping_gen.get_image_embedding_field()
  
b926f678   tangwang   多语言查询
91
92
93
          # Query builder - use multi-language version
          self.query_builder = MultiLanguageQueryBuilder(
              config=config,
be52af70   tangwang   first commit
94
              index_name=config.es_index_name,
be52af70   tangwang   first commit
95
96
97
98
99
100
101
102
103
104
              text_embedding_field=self.text_embedding_field,
              image_embedding_field=self.image_embedding_field
          )
  
      def search(
          self,
          query: str,
          size: int = 10,
          from_: int = 0,
          filters: Optional[Dict[str, Any]] = None,
16c42787   tangwang   feat: implement r...
105
          min_score: Optional[float] = None,
c86c8237   tangwang   支持聚合。过滤项补充了逻辑,但是有问题
106
107
108
109
          context: Optional[RequestContext] = None,
          aggregations: Optional[Dict[str, Any]] = None,
          sort_by: Optional[str] = None,
          sort_order: Optional[str] = "desc"
be52af70   tangwang   first commit
110
111
112
113
114
115
116
117
118
      ) -> SearchResult:
          """
          Execute search query.
  
          Args:
              query: Search query string
              size: Number of results to return
              from_: Offset for pagination
              filters: Additional filters (field: value pairs)
be52af70   tangwang   first commit
119
              min_score: Minimum score threshold
16c42787   tangwang   feat: implement r...
120
              context: Request context for tracking (created if not provided)
c86c8237   tangwang   支持聚合。过滤项补充了逻辑,但是有问题
121
122
123
              aggregations: Aggregation specifications for faceted search
              sort_by: Field name for sorting
              sort_order: Sort order: 'asc' or 'desc'
be52af70   tangwang   first commit
124
125
126
127
  
          Returns:
              SearchResult object
          """
16c42787   tangwang   feat: implement r...
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
          # Create context if not provided (backward compatibility)
          if context is None:
              context = create_request_context()
  
          # Always use config defaults (these are backend configuration, not user parameters)
          enable_translation = self.config.query_config.enable_translation
          enable_embedding = self.config.query_config.enable_text_embedding
          enable_rerank = True  # Always enable reranking as it's part of the search logic
  
          # Start timing
          context.start_stage(RequestContextStage.TOTAL)
  
          context.logger.info(
              f"开始搜索请求 | 查询: '{query}' | 参数: size={size}, from_={from_}, "
              f"enable_translation={enable_translation}, enable_embedding={enable_embedding}, "
              f"enable_rerank={enable_rerank}, min_score={min_score}",
              extra={'reqid': context.reqid, 'uid': context.uid}
          )
  
          # Store search parameters in context
          context.metadata['search_params'] = {
              'size': size,
              'from_': from_,
              'filters': filters,
              'enable_translation': enable_translation,
              'enable_embedding': enable_embedding,
              'enable_rerank': enable_rerank,
c86c8237   tangwang   支持聚合。过滤项补充了逻辑,但是有问题
155
156
157
158
              'min_score': min_score,
              'aggregations': aggregations,
              'sort_by': sort_by,
              'sort_order': sort_order
16c42787   tangwang   feat: implement r...
159
          }
be52af70   tangwang   first commit
160
  
16c42787   tangwang   feat: implement r...
161
162
163
164
165
          context.metadata['feature_flags'] = {
              'translation_enabled': enable_translation,
              'embedding_enabled': enable_embedding,
              'rerank_enabled': enable_rerank
          }
be52af70   tangwang   first commit
166
167
  
          # Step 1: Parse query
16c42787   tangwang   feat: implement r...
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
          context.start_stage(RequestContextStage.QUERY_PARSING)
          try:
              parsed_query = self.query_parser.parse(
                  query,
                  generate_vector=enable_embedding,
                  context=context
              )
              # Store query analysis results in context
              context.store_query_analysis(
                  original_query=parsed_query.original_query,
                  normalized_query=parsed_query.normalized_query,
                  rewritten_query=parsed_query.rewritten_query,
                  detected_language=parsed_query.detected_language,
                  translations=parsed_query.translations,
                  query_vector=parsed_query.query_vector.tolist() if parsed_query.query_vector is not None else None,
                  domain=parsed_query.domain,
                  is_simple_query=self.boolean_parser.is_simple_query(parsed_query.rewritten_query)
              )
be52af70   tangwang   first commit
186
  
16c42787   tangwang   feat: implement r...
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
              context.logger.info(
                  f"查询解析完成 | 原查询: '{parsed_query.original_query}' | "
                  f"重写后: '{parsed_query.rewritten_query}' | "
                  f"语言: {parsed_query.detected_language} | "
                  f"域: {parsed_query.domain} | "
                  f"向量: {'是' if parsed_query.query_vector is not None else '否'}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
          except Exception as e:
              context.set_error(e)
              context.logger.error(
                  f"查询解析失败 | 错误: {str(e)}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
              raise
          finally:
              context.end_stage(RequestContextStage.QUERY_PARSING)
  
          # Step 2: Boolean parsing
          context.start_stage(RequestContextStage.BOOLEAN_PARSING)
          try:
              query_node = None
              if self.boolean_parser.is_simple_query(parsed_query.rewritten_query):
                  # Simple query
                  query_text = parsed_query.rewritten_query
                  context.logger.debug(
                      f"简单查询 | 无布尔表达式",
                      extra={'reqid': context.reqid, 'uid': context.uid}
                  )
              else:
                  # Complex boolean query
                  query_node = self.boolean_parser.parse(parsed_query.rewritten_query)
                  query_text = parsed_query.rewritten_query
                  context.store_intermediate_result('query_node', query_node)
                  context.store_intermediate_result('boolean_ast', str(query_node))
                  context.logger.info(
                      f"布尔表达式解析 | AST: {query_node}",
                      extra={'reqid': context.reqid, 'uid': context.uid}
                  )
          except Exception as e:
              context.set_error(e)
              context.logger.error(
                  f"布尔表达式解析失败 | 错误: {str(e)}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
              raise
          finally:
              context.end_stage(RequestContextStage.BOOLEAN_PARSING)
  
          # Step 3: Query building
          context.start_stage(RequestContextStage.QUERY_BUILDING)
          try:
              es_query = self.query_builder.build_multilang_query(
                  parsed_query=parsed_query,
                  query_vector=parsed_query.query_vector if enable_embedding else None,
                  query_node=query_node,
                  filters=filters,
                  size=size,
                  from_=from_,
                  enable_knn=enable_embedding and parsed_query.query_vector is not None,
                  min_score=min_score
              )
be52af70   tangwang   first commit
249
  
16c42787   tangwang   feat: implement r...
250
251
252
253
254
255
256
257
258
              # Add SPU collapse if configured
              if self.config.spu_config.enabled:
                  es_query = self.query_builder.add_spu_collapse(
                      es_query,
                      self.config.spu_config.spu_field,
                      self.config.spu_config.inner_hits_size
                  )
  
              # Add aggregations for faceted search
c86c8237   tangwang   支持聚合。过滤项补充了逻辑,但是有问题
259
260
261
262
263
              if aggregations:
                  # Use dynamic aggregations from request
                  es_query = self.query_builder.add_dynamic_aggregations(es_query, aggregations)
              elif filters:
                  # Fallback to filter-based aggregations
16c42787   tangwang   feat: implement r...
264
265
266
267
                  agg_fields = [f"{k}_keyword" for k in filters.keys() if f"{k}_keyword" in [f.name for f in self.config.fields]]
                  if agg_fields:
                      es_query = self.query_builder.add_aggregations(es_query, agg_fields)
  
c86c8237   tangwang   支持聚合。过滤项补充了逻辑,但是有问题
268
269
270
271
              # Add sorting if specified
              if sort_by:
                  es_query = self.query_builder.add_sorting(es_query, sort_by, sort_order)
  
16c42787   tangwang   feat: implement r...
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
              # Extract size and from from body for ES client parameters
              body_for_es = {k: v for k, v in es_query.items() if k not in ['size', 'from']}
  
              # Store ES query in context
              context.store_intermediate_result('es_query', es_query)
              context.store_intermediate_result('es_body_for_search', body_for_es)
  
              context.logger.info(
                  f"ES查询构建完成 | 大小: {len(str(es_query))}字符 | "
                  f"KNN: {'是' if enable_embedding and parsed_query.query_vector is not None else '否'} | "
                  f"聚合: {'是' if filters else '否'}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
              context.logger.debug(
                  f"ES查询详情: {es_query}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
          except Exception as e:
              context.set_error(e)
              context.logger.error(
                  f"ES查询构建失败 | 错误: {str(e)}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
              raise
          finally:
              context.end_stage(RequestContextStage.QUERY_BUILDING)
  
          # Step 4: Elasticsearch search
          context.start_stage(RequestContextStage.ELASTICSEARCH_SEARCH)
          try:
              es_response = self.es_client.search(
                  index_name=self.config.es_index_name,
                  body=body_for_es,
                  size=size,
                  from_=from_
be52af70   tangwang   first commit
307
308
              )
  
16c42787   tangwang   feat: implement r...
309
310
              # Store ES response in context
              context.store_intermediate_result('es_response', es_response)
be52af70   tangwang   first commit
311
  
16c42787   tangwang   feat: implement r...
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
              # Extract timing from ES response
              es_took = es_response.get('took', 0)
              context.logger.info(
                  f"ES搜索完成 | 耗时: {es_took}ms | "
                  f"命中数: {es_response.get('hits', {}).get('total', {}).get('value', 0)} | "
                  f"最高分: {es_response.get('hits', {}).get('max_score', 0):.3f}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
          except Exception as e:
              context.set_error(e)
              context.logger.error(
                  f"ES搜索执行失败 | 错误: {str(e)}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
              raise
          finally:
              context.end_stage(RequestContextStage.ELASTICSEARCH_SEARCH)
  
          # Step 5: Result processing
          context.start_stage(RequestContextStage.RESULT_PROCESSING)
          try:
              hits = []
              raw_hits = []
  
              if 'hits' in es_response and 'hits' in es_response['hits']:
                  for hit in es_response['hits']['hits']:
                      raw_hits.append(hit)
  
                      result_doc = {
                          '_id': hit['_id'],
                          '_score': hit['_score'],
                          '_source': hit['_source']
                      }
  
                      # Apply custom ranking if enabled
                      if enable_rerank:
                          base_score = hit['_score']
                          knn_score = None
  
                          # Check if KNN was used
                          if 'knn' in es_query:
                              # KNN score would be in the combined score
                              # For simplicity, extract from score
                              knn_score = base_score * 0.2  # Approximate based on our formula
  
                          custom_score = self.ranking_engine.calculate_score(
                              hit,
                              base_score,
                              knn_score
                          )
                          result_doc['_custom_score'] = custom_score
                          result_doc['_original_score'] = base_score
  
                      hits.append(result_doc)
  
                  # Re-sort by custom score if reranking enabled
be52af70   tangwang   first commit
368
                  if enable_rerank:
16c42787   tangwang   feat: implement r...
369
370
371
372
                      hits.sort(key=lambda x: x.get('_custom_score', x['_score']), reverse=True)
                      context.logger.info(
                          f"重排序完成 | 基于自定义评分表达式",
                          extra={'reqid': context.reqid, 'uid': context.uid}
be52af70   tangwang   first commit
373
                      )
be52af70   tangwang   first commit
374
  
16c42787   tangwang   feat: implement r...
375
376
377
              # Store intermediate results in context
              context.store_intermediate_result('raw_hits', raw_hits)
              context.store_intermediate_result('processed_hits', hits)
be52af70   tangwang   first commit
378
  
16c42787   tangwang   feat: implement r...
379
380
381
382
383
384
              # Extract total and max_score
              total = es_response.get('hits', {}).get('total', {})
              if isinstance(total, dict):
                  total_value = total.get('value', 0)
              else:
                  total_value = total
be52af70   tangwang   first commit
385
  
16c42787   tangwang   feat: implement r...
386
              max_score = es_response.get('hits', {}).get('max_score', 0.0)
be52af70   tangwang   first commit
387
  
16c42787   tangwang   feat: implement r...
388
389
              # Extract aggregations
              aggregations = es_response.get('aggregations', {})
be52af70   tangwang   first commit
390
  
16c42787   tangwang   feat: implement r...
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
              context.logger.info(
                  f"结果处理完成 | 返回: {len(hits)}条 | 总计: {total_value}条 | "
                  f"重排序: {'是' if enable_rerank else '否'}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
  
          except Exception as e:
              context.set_error(e)
              context.logger.error(
                  f"结果处理失败 | 错误: {str(e)}",
                  extra={'reqid': context.reqid, 'uid': context.uid}
              )
              raise
          finally:
              context.end_stage(RequestContextStage.RESULT_PROCESSING)
be52af70   tangwang   first commit
406
  
16c42787   tangwang   feat: implement r...
407
408
409
          # End total timing and build result
          total_duration = context.end_stage(RequestContextStage.TOTAL)
          context.performance_metrics.total_duration = total_duration
be52af70   tangwang   first commit
410
411
412
413
414
415
  
          # Build result
          result = SearchResult(
              hits=hits,
              total=total_value,
              max_score=max_score,
16c42787   tangwang   feat: implement r...
416
              took_ms=int(total_duration),
be52af70   tangwang   first commit
417
418
419
420
              aggregations=aggregations,
              query_info=parsed_query.to_dict()
          )
  
16c42787   tangwang   feat: implement r...
421
422
          # Log complete performance summary
          context.log_performance_summary()
be52af70   tangwang   first commit
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
  
          return result
  
      def search_by_image(
          self,
          image_url: str,
          size: int = 10,
          filters: Optional[Dict[str, Any]] = None
      ) -> SearchResult:
          """
          Search by image similarity.
  
          Args:
              image_url: URL of query image
              size: Number of results
              filters: Additional filters
  
          Returns:
              SearchResult object
          """
          if not self.image_embedding_field:
              raise ValueError("Image embedding field not configured")
  
          # Generate image embedding
          from embeddings import CLIPImageEncoder
          image_encoder = CLIPImageEncoder()
          image_vector = image_encoder.encode_image_from_url(image_url)
  
          if image_vector is None:
              raise ValueError(f"Failed to encode image: {image_url}")
  
          # Build KNN query
          es_query = {
              "size": size,
              "knn": {
                  "field": self.image_embedding_field,
                  "query_vector": image_vector.tolist(),
                  "k": size,
                  "num_candidates": size * 10
              }
          }
  
          if filters:
              es_query["query"] = {
                  "bool": {
                      "filter": self.query_builder._build_filters(filters)
                  }
              }
  
          # Execute search
          es_response = self.es_client.search(
              index_name=self.config.es_index_name,
              body=es_query,
              size=size
          )
  
          # Process results (similar to text search)
          hits = []
          if 'hits' in es_response and 'hits' in es_response['hits']:
              for hit in es_response['hits']['hits']:
                  hits.append({
                      '_id': hit['_id'],
                      '_score': hit['_score'],
                      '_source': hit['_source']
                  })
  
          total = es_response.get('hits', {}).get('total', {})
          if isinstance(total, dict):
              total_value = total.get('value', 0)
          else:
              total_value = total
  
          return SearchResult(
              hits=hits,
              total=total_value,
              max_score=es_response.get('hits', {}).get('max_score', 0.0),
              took_ms=es_response.get('took', 0),
              query_info={'image_url': image_url, 'search_type': 'image_similarity'}
          )
  
b926f678   tangwang   多语言查询
503
504
505
506
507
508
509
510
511
      def get_domain_summary(self) -> Dict[str, Any]:
          """
          Get summary of all configured domains.
  
          Returns:
              Dictionary with domain information
          """
          return self.query_builder.get_domain_summary()
  
be52af70   tangwang   first commit
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
      def get_document(self, doc_id: str) -> Optional[Dict[str, Any]]:
          """
          Get single document by ID.
  
          Args:
              doc_id: Document ID
  
          Returns:
              Document or None if not found
          """
          try:
              response = self.es_client.client.get(
                  index=self.config.es_index_name,
                  id=doc_id
              )
              return response.get('_source')
          except Exception as e:
              print(f"[Searcher] Failed to get document {doc_id}: {e}")
              return None