Blame view

main.py 6.04 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
23
24
25
26
  from config import ConfigLoader
  from utils import ESClient
  from indexer import DataTransformer, IndexingPipeline
  from embeddings import BgeEncoder, CLIPImageEncoder
  from search import Searcher
  
be52af70   tangwang   first commit
27
28
29
  
  def cmd_ingest(args):
      """Run data ingestion."""
4d824a77   tangwang   所有租户共用一套统一配置.tena...
30
      print("Starting data ingestion")
be52af70   tangwang   first commit
31
32
  
      # Load config
4d824a77   tangwang   所有租户共用一套统一配置.tena...
33
34
      config_loader = ConfigLoader("config/config.yaml")
      config = config_loader.load_config()
be52af70   tangwang   first commit
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
  
      # 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
68
69
      os.environ['ES_HOST'] = args.es_host
  
4d824a77   tangwang   所有租户共用一套统一配置.tena...
70
      print("Starting API service (multi-tenant)...")
be52af70   tangwang   first commit
71
72
73
74
75
76
77
78
79
80
81
82
83
      print(f"  Host: {args.host}:{args.port}")
      print(f"  Elasticsearch: {args.es_host}")
  
      uvicorn.run(
          "api.app:app",
          host=args.host,
          port=args.port,
          reload=args.reload
      )
  
  
  def cmd_search(args):
      """Test search from command line."""
be52af70   tangwang   first commit
84
      # Load config
4d824a77   tangwang   所有租户共用一套统一配置.tena...
85
86
      config_loader = ConfigLoader("config/config.yaml")
      config = config_loader.load_config()
be52af70   tangwang   first commit
87
88
89
90
91
92
93
  
      # 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...
94
95
      from query import QueryParser
      query_parser = QueryParser(config)
9f96d6f3   tangwang   短query不用语义搜索
96
      searcher = Searcher(es_client, config, query_parser)
be52af70   tangwang   first commit
97
98
  
      # Execute search
4d824a77   tangwang   所有租户共用一套统一配置.tena...
99
      print(f"Searching for: '{args.query}' (tenant: {args.tenant_id})")
be52af70   tangwang   first commit
100
101
      result = searcher.search(
          query=args.query,
4d824a77   tangwang   所有租户共用一套统一配置.tena...
102
103
          tenant_id=args.tenant_id,
          size=args.size
be52af70   tangwang   first commit
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
      )
  
      # 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
138
139
140
141
142
143
144
      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...
145
      serve_parser = subparsers.add_parser('serve', help='Start API service (multi-tenant)')
be52af70   tangwang   first commit
146
      serve_parser.add_argument('--host', default='0.0.0.0', help='Host to bind to')
2a76641e   tangwang   config
147
      serve_parser.add_argument('--port', type=int, default=6002, help='Port to bind to')
be52af70   tangwang   first commit
148
149
150
151
152
153
      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')
  
      # 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...
154
      search_parser.add_argument('--tenant-id', required=True, help='Tenant ID (required)')
be52af70   tangwang   first commit
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
      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)
      elif args.command == 'search':
          return cmd_search(args)
      else:
          print(f"Unknown command: {args.command}")
          return 1
  
  
  if __name__ == "__main__":
      sys.exit(main())