test_topic_table_query_generation.py 9.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269
  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. "composition_items": [
  51. {
  52. "id": 25684,
  53. "item_level": "element",
  54. "dimension": "形式",
  55. "point_type": "灵感点",
  56. "element_name": "低角度仰拍",
  57. "category_path": "/呈现/视觉/低角度",
  58. "reason": "增强视觉张力",
  59. "is_active": True,
  60. "sources": [
  61. {
  62. "derivation_type": "external_search",
  63. "source_type": "search_case",
  64. "source_reference_id": "case-1",
  65. "reason": "草坪案例证明低角度有动感",
  66. "dataset_from": "",
  67. "source_reference_data": {"large": "must not enter prompt"},
  68. }
  69. ],
  70. },
  71. {
  72. "id": 25688,
  73. "item_level": "element",
  74. "dimension": "实质",
  75. "point_type": "灵感点",
  76. "element_name": "羊",
  77. "reason": "与账号名形成双关",
  78. "is_active": True,
  79. "sources": [],
  80. },
  81. {
  82. "id": 25700,
  83. "item_level": "element",
  84. "dimension": "意图",
  85. "point_type": "目的点",
  86. "element_name": "治愈感",
  87. "reason": "满足粉丝对美好生活的向往",
  88. "is_active": True,
  89. "sources": [],
  90. },
  91. ],
  92. "item_relations": [
  93. {
  94. "source_item_id": 25684,
  95. "target_item_id": 25688,
  96. "reason": "低角度用于突出人与小羊互动",
  97. }
  98. ],
  99. }
  100. ],
  101. }
  102. def _llm_output() -> dict[str, Any]:
  103. fields = [
  104. ("point:8371:point_result", "灵感点怎么拍出来"),
  105. ("item:25684:element_name", "低角度仰拍怎么执行"),
  106. ("item:25688:element_name", "小羊怎么自然入镜"),
  107. ("item:25700:element_name", "治愈感怎么通过画面表现"),
  108. ("topic_direction.persona_dimensions.形式偏好", "怎么保持账号原来的拍法"),
  109. ("item:25684:source:0:reason", "旧案例换场景后怎么调整"),
  110. ]
  111. return {
  112. "field_mappings": [
  113. {
  114. "field_key": field_key,
  115. "judgement": judgement,
  116. "queries": [
  117. f"{judgement}有哪些方法",
  118. "逆光人像怎么拍出彩虹光晕" if index == 0 else f"{judgement}常见失败原因",
  119. ],
  120. }
  121. for index, (field_key, judgement) in enumerate(fields)
  122. ]
  123. }
  124. def test_topic_source_calls_read_only_detail_endpoint():
  125. calls = []
  126. class Response:
  127. def raise_for_status(self):
  128. return None
  129. def json(self):
  130. return _topic_payload()
  131. source = PatternTopicBuildSource(
  132. PatternTopicBuildConfig(base_url="https://pattern.example", timeout_seconds=8),
  133. http_get=lambda url, **kwargs: calls.append((url, kwargs)) or Response(),
  134. )
  135. payload = source.fetch(1229)
  136. assert calls == [
  137. ("https://pattern.example/api/pattern/topic_builds/1229", {"timeout": 8})
  138. ]
  139. assert payload["build"]["id"] == 1229
  140. assert payload["source_url"].endswith("/1229")
  141. def test_topic_table_generation_maps_real_fields_and_uses_unified_budget():
  142. seen = {}
  143. def chat(system, user):
  144. seen["system"] = system
  145. seen["user"] = user
  146. return _llm_output()
  147. preview = generate_topic_table_preview(
  148. _topic_payload(),
  149. topic_id=1392,
  150. max_queries=8,
  151. chat_fn=chat,
  152. )
  153. data = preview_to_dict(preview)
  154. assert data["generator_kind"] == "topic_table"
  155. assert len(data["knowledge_needs"]) == 6
  156. assert set(data["summary"]["route_counts"]) == set(ROUTES)
  157. assert data["summary"]["selected_count"] == 8
  158. assert data["summary"]["persisted"] is False
  159. assert data["summary"]["search_started"] is False
  160. assert data["topic"]["dimensions"] == {
  161. "实质": ["羊"],
  162. "形式": ["低角度仰拍"],
  163. "意图": ["治愈感"],
  164. }
  165. assert "source_reference_data" not in seen["user"]
  166. assert "field_catalog" in seen["user"]
  167. assert "topics[id=1392].composition_items[id=25688].element_name" in seen["user"]
  168. assert "废弃" not in seen["user"]
  169. assert "大模型要判断什么" in seen["system"]
  170. assert data["topic"]["input_groups"]["points"]["灵感点"] == [
  171. "低角度仰拍小羊和彩虹光晕"
  172. ]
  173. assert data["topic"]["input_groups"]["reference_notes"] == [
  174. {"element": "低角度仰拍", "note": "草坪案例证明低角度有动感"}
  175. ]
  176. assert data["field_mappings"][0] == {
  177. "field_key": "point:8371:point_result",
  178. "field_path": "topics[id=1392].points[id=8371].point_result",
  179. "field_label": "灵感点字段",
  180. "field_value": "低角度仰拍小羊和彩虹光晕",
  181. "judgement": "灵感点怎么拍出来",
  182. "queries": [
  183. "灵感点怎么拍出来有哪些方法",
  184. "逆光人像怎么拍出彩虹光晕",
  185. ],
  186. }
  187. def test_topic_table_generation_rejects_unknown_field_key():
  188. malformed = _llm_output()
  189. malformed["field_mappings"][-1]["field_key"] = "invented-field"
  190. with pytest.raises(TopicTableGenerationError, match="不存在的选题表字段"):
  191. generate_topic_table_preview(
  192. _topic_payload(),
  193. topic_id=1392,
  194. max_queries=18,
  195. chat_fn=lambda _system, _user: malformed,
  196. )
  197. def test_topic_table_generation_rejects_duplicate_field_key():
  198. malformed = _llm_output()
  199. malformed["field_mappings"][-1]["field_key"] = malformed["field_mappings"][0][
  200. "field_key"
  201. ]
  202. with pytest.raises(TopicTableGenerationError, match="重复引用"):
  203. generate_topic_table_preview(
  204. _topic_payload(),
  205. topic_id=1392,
  206. max_queries=18,
  207. chat_fn=lambda _system, _user: malformed,
  208. )
  209. def test_topic_table_prompt_and_preview_api_do_not_persist_or_search(monkeypatch):
  210. class Source:
  211. def fetch(self, topic_build_id):
  212. assert topic_build_id == 1229
  213. return _topic_payload()
  214. monkeypatch.setattr(query_generation, "_topic_source", lambda _env: Source())
  215. monkeypatch.setattr(
  216. query_generation,
  217. "_topic_query_chat",
  218. lambda system, user, settings: _llm_output(),
  219. )
  220. monkeypatch.setattr(
  221. query_generation.Settings,
  222. "from_env",
  223. classmethod(lambda cls, _env: object()),
  224. )
  225. client = TestClient(app)
  226. prompt = client.get("/api/query-generation/topic-table/prompt")
  227. assert prompt.status_code == 200
  228. assert prompt.json()["example_request"]["topic_build_id"] == 1229
  229. assert len(prompt.json()["routes"]) == 6
  230. response = client.post(
  231. "/api/query-generation/topic-table/preview",
  232. json={"topic_build_id": 1229, "topic_id": 1392, "max_queries": 10},
  233. )
  234. assert response.status_code == 200
  235. body = response.json()
  236. assert body["source"]["topic_id"] == 1392
  237. assert body["summary"]["selected_count"] == 10
  238. assert body["summary"]["persisted"] is False
  239. assert body["summary"]["search_started"] is False