Blame view

tests/test_llm_enrichment_batch_fill.py 2.33 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
                  "id": item["id"],
                  "qanchors": {
                      "zh": [f"zh-anchor-{item['id']}"],
                      "en": [f"en-anchor-{item['id']}"],
                  },
36516857   tangwang   feat(product_enri...
22
                  "enriched_tags": {"zh": ["t1", "t2"], "en": ["t1", "t2"]},
d350861f   tangwang   索引结构修改
23
                  "enriched_attributes": [
80f1e036   tangwang   enriched_attribut...
24
                      {"name": "tags", "value": {"zh": ["t1"], "en": ["t1"]}},
d350861f   tangwang   索引结构修改
25
                  ],
36516857   tangwang   feat(product_enri...
26
27
28
                  "enriched_taxonomy_attributes": [
                      {"name": "Product Type", "value": {"zh": ["连衣裙"], "en": ["dress"]}},
                  ],
be3f0d46   tangwang   /indexer/enrich-c...
29
              }
d350861f   tangwang   索引结构修改
30
              for item in items
be3f0d46   tangwang   /indexer/enrich-c...
31
32
33
34
          ]
  
      import indexer.document_transformer as doc_tr
  
d350861f   tangwang   索引结构修改
35
      monkeypatch.setattr(doc_tr, "build_index_content_fields", _fake_build_index_content_fields)
be3f0d46   tangwang   /indexer/enrich-c...
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
  
      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   索引结构修改
56
      assert seen_calls == [{"n": 45, "tenant_id": "162"}]
be3f0d46   tangwang   /indexer/enrich-c...
57
  
d350861f   tangwang   索引结构修改
58
59
      assert docs[0]["qanchors"]["zh"] == ["zh-anchor-0"]
      assert docs[0]["qanchors"]["en"] == ["en-anchor-0"]
36516857   tangwang   feat(product_enri...
60
61
      assert docs[0]["enriched_tags"]["zh"] == ["t1", "t2"]
      assert docs[0]["enriched_tags"]["en"] == ["t1", "t2"]
80f1e036   tangwang   enriched_attribut...
62
      assert {"name": "tags", "value": {"zh": ["t1"], "en": ["t1"]}} in docs[0]["enriched_attributes"]
36516857   tangwang   feat(product_enri...
63
64
65
66
      assert {
          "name": "Product Type",
          "value": {"zh": ["连衣裙"], "en": ["dress"]},
      } in docs[0]["enriched_taxonomy_attributes"]