test_topic_table_query_generation.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332
  1. from __future__ import annotations
  2. from typing import Any
  3. import pytest
  4. from fastapi.testclient import TestClient
  5. from app.api import app
  6. from app.routes import query_generation
  7. from query_planning.sources.pattern_topic_build import (
  8. PatternTopicBuildConfig,
  9. PatternTopicBuildSource,
  10. )
  11. from query_planning.topic_table import (
  12. ROUTES,
  13. TopicTableGenerationError,
  14. generate_topic_table_preview,
  15. preview_to_dict,
  16. )
  17. def _topic_payload() -> dict[str, Any]:
  18. return {
  19. "success": True,
  20. "source_url": "https://pattern.aiddit.com/api/pattern/topic_builds/1229",
  21. "build": {
  22. "id": 1229,
  23. "demand": "为小羊糊涂啊创作新选题",
  24. "demand_constraints": None,
  25. "strategies_config": {"always_on": []},
  26. "status": "success",
  27. },
  28. "topics": [
  29. {
  30. "id": 1392,
  31. "status": "mature",
  32. "result": "樱花树下低角度拍少女与小羊,传递治愈感并分享穿搭。",
  33. "topic_direction": {
  34. "account_name": "小羊糊涂啊",
  35. "persona_dimensions": {
  36. "实质偏好": "樱花与穿搭",
  37. "形式偏好": "低角度和动态摆拍",
  38. "意图偏好": "治愈和种草",
  39. },
  40. },
  41. "points": [
  42. {
  43. "id": 8371,
  44. "point_type": "灵感点",
  45. "point_result": "低角度仰拍小羊和彩虹光晕",
  46. "is_active": True,
  47. "item_ids": [25684, 25688],
  48. },
  49. {
  50. "id": 8372,
  51. "point_type": "目的点",
  52. "point_result": "传递治愈感",
  53. "is_active": True,
  54. "item_ids": [25700],
  55. },
  56. {
  57. "id": 8374,
  58. "point_type": "关键点",
  59. "point_result": "逆光氛围感特写",
  60. "is_active": True,
  61. "item_ids": [25684],
  62. },
  63. ],
  64. "composition_items": [
  65. {
  66. "id": 25684,
  67. "item_level": "element",
  68. "dimension": "形式",
  69. "point_type": "灵感点",
  70. "element_name": "低角度仰拍",
  71. "category_path": "/呈现/视觉/低角度",
  72. "reason": "增强视觉张力",
  73. "is_active": True,
  74. "sources": [
  75. {
  76. "derivation_type": "external_search",
  77. "source_type": "search_case",
  78. "source_reference_id": "case-1",
  79. "reason": "草坪案例证明低角度有动感",
  80. "dataset_from": "",
  81. "source_reference_data": {"large": "must not enter prompt"},
  82. }
  83. ],
  84. },
  85. {
  86. "id": 25688,
  87. "item_level": "element",
  88. "dimension": "实质",
  89. "point_type": "灵感点",
  90. "element_name": "羊",
  91. "reason": "与账号名形成双关",
  92. "is_active": True,
  93. "sources": [],
  94. },
  95. {
  96. "id": 25700,
  97. "item_level": "element",
  98. "dimension": "意图",
  99. "point_type": "目的点",
  100. "element_name": "治愈感",
  101. "reason": "满足粉丝对美好生活的向往",
  102. "is_active": True,
  103. "sources": [],
  104. },
  105. ],
  106. "item_relations": [
  107. {
  108. "source_item_id": 25684,
  109. "target_item_id": 25688,
  110. "reason": "低角度用于突出人与小羊互动",
  111. }
  112. ],
  113. }
  114. ],
  115. }
  116. def _llm_output() -> dict[str, Any]:
  117. fields = [
  118. ("point:8371:point_result", "灵感点怎么拍出来"),
  119. ("point:8372:point_result", "治愈感怎么落实到画面"),
  120. ("point:8374:point_result", "逆光特写怎么执行"),
  121. ("item:25688:element_name", "小羊怎么自然入镜"),
  122. ("item:25684:element_name", "低角度仰拍怎么执行"),
  123. ("item:25700:element_name", "治愈感怎么通过画面表现"),
  124. ("topic_direction.persona_dimensions.实质偏好", "怎么保持账号原来的选材"),
  125. ("topic_direction.persona_dimensions.形式偏好", "怎么保持账号原来的拍法"),
  126. ("topic_direction.persona_dimensions.意图偏好", "怎么保持账号原来的表达效果"),
  127. ("item:25684:source:0:reason", "旧案例换场景后怎么调整"),
  128. ]
  129. return {
  130. "field_mappings": [
  131. {
  132. "field_key": field_key,
  133. "judgement": judgement,
  134. "queries": [
  135. f"{judgement}有哪些方法",
  136. "逆光人像怎么拍出彩虹光晕" if index == 0 else f"{judgement}常见失败原因",
  137. ],
  138. }
  139. for index, (field_key, judgement) in enumerate(fields)
  140. ]
  141. }
  142. def test_topic_source_calls_read_only_detail_endpoint():
  143. calls = []
  144. class Response:
  145. def raise_for_status(self):
  146. return None
  147. def json(self):
  148. return _topic_payload()
  149. source = PatternTopicBuildSource(
  150. PatternTopicBuildConfig(base_url="https://pattern.example", timeout_seconds=8),
  151. http_get=lambda url, **kwargs: calls.append((url, kwargs)) or Response(),
  152. )
  153. payload = source.fetch(1229)
  154. assert calls == [
  155. ("https://pattern.example/api/pattern/topic_builds/1229", {"timeout": 8})
  156. ]
  157. assert payload["build"]["id"] == 1229
  158. assert payload["source_url"].endswith("/1229")
  159. def test_topic_table_generation_maps_real_fields_and_uses_unified_budget():
  160. seen = {}
  161. def chat(system, user):
  162. seen["system"] = system
  163. seen["user"] = user
  164. return _llm_output()
  165. preview = generate_topic_table_preview(
  166. _topic_payload(),
  167. topic_id=1392,
  168. max_queries=8,
  169. chat_fn=chat,
  170. )
  171. data = preview_to_dict(preview)
  172. assert data["generator_kind"] == "topic_table"
  173. assert len(data["knowledge_needs"]) == 10
  174. assert set(data["summary"]["route_counts"]) == set(ROUTES)
  175. assert data["summary"]["selected_count"] == 8
  176. assert data["summary"]["persisted"] is False
  177. assert data["summary"]["search_started"] is False
  178. assert len(data["source_groups"]) == 6
  179. assert all(group["field_mappings"] for group in data["source_groups"])
  180. assert [group["field_count"] for group in data["source_groups"]] == [1, 1, 1, 5, 3, 1]
  181. assert sum(len(group["fields"]) for group in data["source_groups"]) == 12
  182. topic_result_field = next(
  183. field
  184. for field in data["source_groups"][3]["fields"]
  185. if field["field_key"] == "topic.result"
  186. )
  187. assert topic_result_field["used_for_query"] is False
  188. assert topic_result_field["judgement"] is None
  189. assert topic_result_field["queries"] == []
  190. assert data["topic"]["dimensions"] == {
  191. "实质": ["羊"],
  192. "形式": ["低角度仰拍"],
  193. "意图": ["治愈感"],
  194. }
  195. assert "source_reference_data" not in seen["user"]
  196. assert "field_catalog" in seen["user"]
  197. assert "topics[id=1392].composition_items[id=25688].element_name" in seen["user"]
  198. assert "废弃" not in seen["user"]
  199. assert "大模型要判断什么" in seen["system"]
  200. assert data["topic"]["input_groups"]["points"]["灵感点"] == [
  201. "低角度仰拍小羊和彩虹光晕"
  202. ]
  203. assert data["topic"]["input_groups"]["reference_notes"] == [
  204. {"element": "低角度仰拍", "note": "草坪案例证明低角度有动感"}
  205. ]
  206. assert data["field_mappings"][0] == {
  207. "field_key": "point:8371:point_result",
  208. "field_path": "topics[id=1392].points[id=8371].point_result",
  209. "field_label": "灵感点字段",
  210. "field_value": "低角度仰拍小羊和彩虹光晕",
  211. "judgement": "灵感点怎么拍出来",
  212. "source_group_key": "unique_hook",
  213. "source_group_label": "灵感点",
  214. "queries": [
  215. "灵感点怎么拍出来有哪些方法",
  216. "逆光人像怎么拍出彩虹光晕",
  217. ],
  218. }
  219. def test_topic_table_generation_rejects_unknown_field_key():
  220. malformed = _llm_output()
  221. malformed["field_mappings"][-1]["field_key"] = "invented-field"
  222. with pytest.raises(TopicTableGenerationError, match="不存在的选题表字段"):
  223. generate_topic_table_preview(
  224. _topic_payload(),
  225. topic_id=1392,
  226. max_queries=18,
  227. chat_fn=lambda _system, _user: malformed,
  228. )
  229. def test_topic_table_generation_rejects_duplicate_field_key():
  230. malformed = _llm_output()
  231. malformed["field_mappings"][-1]["field_key"] = malformed["field_mappings"][0][
  232. "field_key"
  233. ]
  234. with pytest.raises(TopicTableGenerationError, match="重复引用"):
  235. generate_topic_table_preview(
  236. _topic_payload(),
  237. topic_id=1392,
  238. max_queries=18,
  239. chat_fn=lambda _system, _user: malformed,
  240. )
  241. def test_topic_table_generation_rejects_missing_source_group():
  242. malformed = _llm_output()
  243. malformed["field_mappings"] = malformed["field_mappings"][:-1]
  244. with pytest.raises(TopicTableGenerationError, match="全部六种数据来源"):
  245. generate_topic_table_preview(
  246. _topic_payload(),
  247. topic_id=1392,
  248. max_queries=18,
  249. chat_fn=lambda _system, _user: malformed,
  250. )
  251. def test_topic_table_generation_retries_once_when_a_source_group_is_missing():
  252. malformed = _llm_output()
  253. malformed["field_mappings"] = malformed["field_mappings"][:-1]
  254. outputs = iter([malformed, _llm_output()])
  255. prompts = []
  256. preview = generate_topic_table_preview(
  257. _topic_payload(),
  258. topic_id=1392,
  259. max_queries=18,
  260. chat_fn=lambda _system, user: prompts.append(user) or next(outputs),
  261. )
  262. assert len(preview.result.selected_queries) == 18
  263. assert len(prompts) == 2
  264. assert "上一次输出没有满足字段覆盖要求" in prompts[1]
  265. def test_topic_table_prompt_and_preview_api_do_not_persist_or_search(monkeypatch):
  266. class Source:
  267. def fetch(self, topic_build_id):
  268. assert topic_build_id == 1229
  269. return _topic_payload()
  270. monkeypatch.setattr(query_generation, "_topic_source", lambda _env: Source())
  271. monkeypatch.setattr(
  272. query_generation,
  273. "_topic_query_chat",
  274. lambda system, user, settings: _llm_output(),
  275. )
  276. monkeypatch.setattr(
  277. query_generation.Settings,
  278. "from_env",
  279. classmethod(lambda cls, _env: object()),
  280. )
  281. client = TestClient(app)
  282. prompt = client.get("/api/query-generation/topic-table/prompt")
  283. assert prompt.status_code == 200
  284. assert prompt.json()["example_request"]["topic_build_id"] == 1229
  285. assert len(prompt.json()["routes"]) == 6
  286. response = client.post(
  287. "/api/query-generation/topic-table/preview",
  288. json={"topic_build_id": 1229, "topic_id": 1392, "max_queries": 10},
  289. )
  290. assert response.status_code == 200
  291. body = response.json()
  292. assert body["source"]["topic_id"] == 1392
  293. assert body["summary"]["selected_count"] == 10
  294. assert body["summary"]["persisted"] is False
  295. assert body["summary"]["search_started"] is False