Blame view

tests/test_llm_enrichment_batch_fill.py 2.1 KB
be3f0d46   tangwang   /indexer/enrich-c...
1
2
3
4
5
6
7
8
9
  from __future__ import annotations
  
  from typing import Any, Dict, List
  
  import pandas as pd
  
  from indexer.document_transformer import SPUDocumentTransformer
  
  
d350861f   tangwang   索引结构修改
10
  def test_fill_llm_attributes_batch_uses_product_enrich_helper(monkeypatch):
be3f0d46   tangwang   /indexer/enrich-c...
11
12
      seen_calls: List[Dict[str, Any]] = []
  
d350861f   tangwang   索引结构修改
13
14
      def _fake_build_index_content_fields(items, tenant_id=None):
          seen_calls.append({"n": len(items), "tenant_id": tenant_id})
be3f0d46   tangwang   /indexer/enrich-c...
15
16
          return [
              {
d350861f   tangwang   索引结构修改
17
18
19
20
21
22
23
24
25
26
                  "id": item["id"],
                  "qanchors": {
                      "zh": [f"zh-anchor-{item['id']}"],
                      "en": [f"en-anchor-{item['id']}"],
                  },
                  "tags": {"zh": ["t1", "t2"], "en": ["t1", "t2"]},
                  "enriched_attributes": [
                      {"name": "tags", "value": {"zh": "t1"}},
                      {"name": "tags", "value": {"en": "t1"}},
                  ],
be3f0d46   tangwang   /indexer/enrich-c...
27
              }
d350861f   tangwang   索引结构修改
28
              for item in items
be3f0d46   tangwang   /indexer/enrich-c...
29
30
31
32
          ]
  
      import indexer.document_transformer as doc_tr
  
d350861f   tangwang   索引结构修改
33
      monkeypatch.setattr(doc_tr, "build_index_content_fields", _fake_build_index_content_fields)
be3f0d46   tangwang   /indexer/enrich-c...
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
  
      transformer = SPUDocumentTransformer(
          category_id_to_name={},
          searchable_option_dimensions=[],
          tenant_config={"index_languages": ["zh", "en"], "primary_language": "zh"},
          translator=None,
          encoder=None,
          enable_title_embedding=False,
          image_encoder=None,
          enable_image_embedding=False,
      )
  
      docs: List[Dict[str, Any]] = []
      rows: List[pd.Series] = []
      for i in range(45):
          docs.append({"tenant_id": "162", "spu_id": str(i)})
          rows.append(pd.Series({"id": i, "title": f"title-{i}"}))
  
      transformer.fill_llm_attributes_batch(docs, rows)
  
d350861f   tangwang   索引结构修改
54
      assert seen_calls == [{"n": 45, "tenant_id": "162"}]
be3f0d46   tangwang   /indexer/enrich-c...
55
  
d350861f   tangwang   索引结构修改
56
57
58
59
60
61
      assert docs[0]["qanchors"]["zh"] == ["zh-anchor-0"]
      assert docs[0]["qanchors"]["en"] == ["en-anchor-0"]
      assert docs[0]["tags"]["zh"] == ["t1", "t2"]
      assert docs[0]["tags"]["en"] == ["t1", "t2"]
      assert {"name": "tags", "value": {"zh": "t1"}} in docs[0]["enriched_attributes"]
      assert {"name": "tags", "value": {"en": "t1"}} in docs[0]["enriched_attributes"]