Blame view

main.py 7.57 KB
be52af70   tangwang   first commit
1
2
3
4
5
6
7
8
9
10
11
12
13
14
  #!/usr/bin/env python3
  """
  Main entry point for SearchEngine operations.
  
  Provides a unified CLI for common operations:
  - ingest: Ingest data into Elasticsearch
  - serve: Start API service
  - search: Test search from command line
  """
  
  import sys
  import os
  import argparse
  import json
a77693fe   tangwang   调整配置目录结构
15
16
  import pandas as pd
  import uvicorn
be52af70   tangwang   first commit
17
18
19
20
  
  # Add parent directory to path
  sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
  
a77693fe   tangwang   调整配置目录结构
21
22
  from config import ConfigLoader
  from utils import ESClient
a77693fe   tangwang   调整配置目录结构
23
24
  from search import Searcher
  
be52af70   tangwang   first commit
25
26
27
  
  def cmd_ingest(args):
      """Run data ingestion."""
bb9c626c   tangwang   搜索服务(6002)不再初始化/挂...
28
29
30
31
32
33
34
35
36
      # Local imports to avoid hard dependency at module import time
      import pandas as pd
      from embeddings import BgeEncoder, CLIPImageEncoder
      from indexer.bulk_indexer import IndexingPipeline
      # NOTE: DataTransformer was referenced historically, but the concrete
      # implementation is now provided via customer-specific scripts
      # (e.g. data/customer1/ingest_customer1.py). If you still need a generic
      # ingestion pipeline here, you can wire your own transformer.
      from indexer.spu_transformer import SPUTransformer as DataTransformer
4d824a77   tangwang   所有租户共用一套统一配置.tena...
37
      print("Starting data ingestion")
be52af70   tangwang   first commit
38
39
  
      # Load config
4d824a77   tangwang   所有租户共用一套统一配置.tena...
40
41
      config_loader = ConfigLoader("config/config.yaml")
      config = config_loader.load_config()
be52af70   tangwang   first commit
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
  
      # Initialize ES
      es_client = ESClient(hosts=[args.es_host])
      if not es_client.ping():
          print(f"ERROR: Cannot connect to Elasticsearch at {args.es_host}")
          return 1
  
      # Load data
      df = pd.read_csv(args.csv_file)
      if args.limit:
          df = df.head(args.limit)
      print(f"Loaded {len(df)} documents")
  
      # Initialize encoders
      text_encoder = None if args.skip_embeddings else BgeEncoder()
      image_encoder = None if args.skip_embeddings else CLIPImageEncoder()
  
      # Transform and index
      transformer = DataTransformer(config, text_encoder, image_encoder, use_cache=True)
      pipeline = IndexingPipeline(config, es_client, transformer, recreate_index=args.recreate)
  
      results = pipeline.run(df, batch_size=args.batch_size)
  
      print(f"\nIngestion complete:")
      print(f"  Success: {results['success']}")
      print(f"  Failed: {results['failed']}")
      print(f"  Time: {results['elapsed_time']:.2f}s")
  
      return 0
  
  
  def cmd_serve(args):
      """Start API service."""
be52af70   tangwang   first commit
75
76
      os.environ['ES_HOST'] = args.es_host
  
4d824a77   tangwang   所有租户共用一套统一配置.tena...
77
      print("Starting API service (multi-tenant)...")
bb9c626c   tangwang   搜索服务(6002)不再初始化/挂...
78
      print(f"  Host: {args.host}:{args.port} (search + indexer routes)")
be52af70   tangwang   first commit
79
80
81
82
83
84
85
86
87
88
      print(f"  Elasticsearch: {args.es_host}")
  
      uvicorn.run(
          "api.app:app",
          host=args.host,
          port=args.port,
          reload=args.reload
      )
  
  
bb9c626c   tangwang   搜索服务(6002)不再初始化/挂...
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
  def cmd_serve_indexer(args):
      """Start dedicated Indexer API service (no search endpoints)."""
      os.environ['ES_HOST'] = args.es_host
  
      print("Starting Indexer API service...")
      print(f"  Host: {args.host}:{args.port} (indexer only)")
      print(f"  Elasticsearch: {args.es_host}")
  
      uvicorn.run(
          "api.indexer_app:app",
          host=args.host,
          port=args.port,
          reload=args.reload
      )
  
be52af70   tangwang   first commit
104
105
  def cmd_search(args):
      """Test search from command line."""
be52af70   tangwang   first commit
106
      # Load config
4d824a77   tangwang   所有租户共用一套统一配置.tena...
107
108
      config_loader = ConfigLoader("config/config.yaml")
      config = config_loader.load_config()
be52af70   tangwang   first commit
109
110
111
112
113
114
115
  
      # Initialize ES and searcher
      es_client = ESClient(hosts=[args.es_host])
      if not es_client.ping():
          print(f"ERROR: Cannot connect to Elasticsearch at {args.es_host}")
          return 1
  
4d824a77   tangwang   所有租户共用一套统一配置.tena...
116
117
      from query import QueryParser
      query_parser = QueryParser(config)
9f96d6f3   tangwang   短query不用语义搜索
118
      searcher = Searcher(es_client, config, query_parser)
be52af70   tangwang   first commit
119
120
  
      # Execute search
4d824a77   tangwang   所有租户共用一套统一配置.tena...
121
      print(f"Searching for: '{args.query}' (tenant: {args.tenant_id})")
be52af70   tangwang   first commit
122
123
      result = searcher.search(
          query=args.query,
4d824a77   tangwang   所有租户共用一套统一配置.tena...
124
125
          tenant_id=args.tenant_id,
          size=args.size
be52af70   tangwang   first commit
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
      )
  
      # Display results
      print(f"\nFound {result.total} results in {result.took_ms}ms")
      print(f"Max score: {result.max_score}")
  
      if args.json:
          print(json.dumps(result.to_dict(), indent=2, ensure_ascii=False))
      else:
          print(f"\nTop {len(result.hits)} results:")
          for i, hit in enumerate(result.hits, 1):
              source = hit['_source']
              score = hit['_score']
              print(f"\n{i}. Score: {score:.4f}")
              print(f"   ID: {source.get('skuId', 'N/A')}")
              print(f"   Name: {source.get('name', 'N/A')}")
              print(f"   Category: {source.get('categoryName', 'N/A')}")
              print(f"   Brand: {source.get('brandName', 'N/A')}")
  
      return 0
  
  
  def main():
      """Main CLI entry point."""
      parser = argparse.ArgumentParser(
          description='SearchEngine - E-Commerce Search SaaS',
          formatter_class=argparse.RawDescriptionHelpFormatter
      )
  
      subparsers = parser.add_subparsers(dest='command', help='Command to execute')
  
      # Ingest command
      ingest_parser = subparsers.add_parser('ingest', help='Ingest data into Elasticsearch')
      ingest_parser.add_argument('csv_file', help='Path to CSV data file')
be52af70   tangwang   first commit
160
161
162
163
164
165
166
      ingest_parser.add_argument('--es-host', default='http://localhost:9200', help='Elasticsearch host')
      ingest_parser.add_argument('--limit', type=int, help='Limit number of documents')
      ingest_parser.add_argument('--batch-size', type=int, default=100, help='Batch size')
      ingest_parser.add_argument('--recreate', action='store_true', help='Recreate index')
      ingest_parser.add_argument('--skip-embeddings', action='store_true', help='Skip embeddings')
  
      # Serve command
4d824a77   tangwang   所有租户共用一套统一配置.tena...
167
      serve_parser = subparsers.add_parser('serve', help='Start API service (multi-tenant)')
be52af70   tangwang   first commit
168
      serve_parser.add_argument('--host', default='0.0.0.0', help='Host to bind to')
2a76641e   tangwang   config
169
      serve_parser.add_argument('--port', type=int, default=6002, help='Port to bind to')
be52af70   tangwang   first commit
170
171
172
      serve_parser.add_argument('--es-host', default='http://localhost:9200', help='Elasticsearch host')
      serve_parser.add_argument('--reload', action='store_true', help='Enable auto-reload')
  
bb9c626c   tangwang   搜索服务(6002)不再初始化/挂...
173
174
175
176
177
178
179
180
181
182
      # Serve-indexer command
      serve_indexer_parser = subparsers.add_parser(
          'serve-indexer',
          help='Start dedicated Indexer API service (indexer routes only)'
      )
      serve_indexer_parser.add_argument('--host', default='0.0.0.0', help='Host to bind to')
      serve_indexer_parser.add_argument('--port', type=int, default=6004, help='Port to bind to')
      serve_indexer_parser.add_argument('--es-host', default='http://localhost:9200', help='Elasticsearch host')
      serve_indexer_parser.add_argument('--reload', action='store_true', help='Enable auto-reload')
  
be52af70   tangwang   first commit
183
184
185
      # Search command
      search_parser = subparsers.add_parser('search', help='Test search from command line')
      search_parser.add_argument('query', help='Search query')
4d824a77   tangwang   所有租户共用一套统一配置.tena...
186
      search_parser.add_argument('--tenant-id', required=True, help='Tenant ID (required)')
be52af70   tangwang   first commit
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
      search_parser.add_argument('--es-host', default='http://localhost:9200', help='Elasticsearch host')
      search_parser.add_argument('--size', type=int, default=10, help='Number of results')
      search_parser.add_argument('--no-translation', action='store_true', help='Disable translation')
      search_parser.add_argument('--no-embedding', action='store_true', help='Disable embeddings')
      search_parser.add_argument('--json', action='store_true', help='Output JSON')
  
      args = parser.parse_args()
  
      if not args.command:
          parser.print_help()
          return 1
  
      # Execute command
      if args.command == 'ingest':
          return cmd_ingest(args)
      elif args.command == 'serve':
          return cmd_serve(args)
bb9c626c   tangwang   搜索服务(6002)不再初始化/挂...
204
205
      elif args.command == 'serve-indexer':
          return cmd_serve_indexer(args)
be52af70   tangwang   first commit
206
207
208
209
210
211
212
213
214
      elif args.command == 'search':
          return cmd_search(args)
      else:
          print(f"Unknown command: {args.command}")
          return 1
  
  
  if __name__ == "__main__":
      sys.exit(main())