test_topic_table_query_generation.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322
  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 data["topic"]["dimensions"] == {
  181. "实质": ["羊"],
  182. "形式": ["低角度仰拍"],
  183. "意图": ["治愈感"],
  184. }
  185. assert "source_reference_data" not in seen["user"]
  186. assert "field_catalog" in seen["user"]
  187. assert "topics[id=1392].composition_items[id=25688].element_name" in seen["user"]
  188. assert "废弃" not in seen["user"]
  189. assert "大模型要判断什么" in seen["system"]
  190. assert data["topic"]["input_groups"]["points"]["灵感点"] == [
  191. "低角度仰拍小羊和彩虹光晕"
  192. ]
  193. assert data["topic"]["input_groups"]["reference_notes"] == [
  194. {"element": "低角度仰拍", "note": "草坪案例证明低角度有动感"}
  195. ]
  196. assert data["field_mappings"][0] == {
  197. "field_key": "point:8371:point_result",
  198. "field_path": "topics[id=1392].points[id=8371].point_result",
  199. "field_label": "灵感点字段",
  200. "field_value": "低角度仰拍小羊和彩虹光晕",
  201. "judgement": "灵感点怎么拍出来",
  202. "source_group_key": "unique_hook",
  203. "source_group_label": "灵感点",
  204. "queries": [
  205. "灵感点怎么拍出来有哪些方法",
  206. "逆光人像怎么拍出彩虹光晕",
  207. ],
  208. }
  209. def test_topic_table_generation_rejects_unknown_field_key():
  210. malformed = _llm_output()
  211. malformed["field_mappings"][-1]["field_key"] = "invented-field"
  212. with pytest.raises(TopicTableGenerationError, match="不存在的选题表字段"):
  213. generate_topic_table_preview(
  214. _topic_payload(),
  215. topic_id=1392,
  216. max_queries=18,
  217. chat_fn=lambda _system, _user: malformed,
  218. )
  219. def test_topic_table_generation_rejects_duplicate_field_key():
  220. malformed = _llm_output()
  221. malformed["field_mappings"][-1]["field_key"] = malformed["field_mappings"][0][
  222. "field_key"
  223. ]
  224. with pytest.raises(TopicTableGenerationError, match="重复引用"):
  225. generate_topic_table_preview(
  226. _topic_payload(),
  227. topic_id=1392,
  228. max_queries=18,
  229. chat_fn=lambda _system, _user: malformed,
  230. )
  231. def test_topic_table_generation_rejects_missing_source_group():
  232. malformed = _llm_output()
  233. malformed["field_mappings"] = malformed["field_mappings"][:-1]
  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_retries_once_when_a_source_group_is_missing():
  242. malformed = _llm_output()
  243. malformed["field_mappings"] = malformed["field_mappings"][:-1]
  244. outputs = iter([malformed, _llm_output()])
  245. prompts = []
  246. preview = generate_topic_table_preview(
  247. _topic_payload(),
  248. topic_id=1392,
  249. max_queries=18,
  250. chat_fn=lambda _system, user: prompts.append(user) or next(outputs),
  251. )
  252. assert len(preview.result.selected_queries) == 18
  253. assert len(prompts) == 2
  254. assert "上一次输出没有满足字段覆盖要求" in prompts[1]
  255. def test_topic_table_prompt_and_preview_api_do_not_persist_or_search(monkeypatch):
  256. class Source:
  257. def fetch(self, topic_build_id):
  258. assert topic_build_id == 1229
  259. return _topic_payload()
  260. monkeypatch.setattr(query_generation, "_topic_source", lambda _env: Source())
  261. monkeypatch.setattr(
  262. query_generation,
  263. "_topic_query_chat",
  264. lambda system, user, settings: _llm_output(),
  265. )
  266. monkeypatch.setattr(
  267. query_generation.Settings,
  268. "from_env",
  269. classmethod(lambda cls, _env: object()),
  270. )
  271. client = TestClient(app)
  272. prompt = client.get("/api/query-generation/topic-table/prompt")
  273. assert prompt.status_code == 200
  274. assert prompt.json()["example_request"]["topic_build_id"] == 1229
  275. assert len(prompt.json()["routes"]) == 6
  276. response = client.post(
  277. "/api/query-generation/topic-table/preview",
  278. json={"topic_build_id": 1229, "topic_id": 1392, "max_queries": 10},
  279. )
  280. assert response.status_code == 200
  281. body = response.json()
  282. assert body["source"]["topic_id"] == 1392
  283. assert body["summary"]["selected_count"] == 10
  284. assert body["summary"]["persisted"] is False
  285. assert body["summary"]["search_started"] is False