test_retrieval_adapters.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317
  1. from __future__ import annotations
  2. import json
  3. from types import SimpleNamespace
  4. import httpx
  5. import pytest
  6. from script_build_host.adapters.retrieval import (
  7. ExternalRetrievalAdapter,
  8. FileDecodeRetrievalAdapter,
  9. FileKnowledgeRetrievalAdapter,
  10. HttpDecodeRetrievalAdapter,
  11. PatternRetrievalAdapter,
  12. SafeHttpClient,
  13. SafeImageAdapter,
  14. )
  15. from script_build_host.adapters.uploaded_topic import SqlUploadedTopicGateway
  16. from script_build_host.domain.errors import ProtocolViolation, UnsafeOutboundTarget
  17. from script_build_host.infrastructure.outbound import OutboundPolicy
  18. from script_build_host.repositories.legacy_input import LegacySqlAlchemyInputReader
  19. async def _public_resolver(_: str, __: int) -> tuple[str, ...]:
  20. return ("93.184.216.34",)
  21. class _RawArtifacts:
  22. def __init__(self) -> None:
  23. self.values: list[bytes] = []
  24. async def freeze_bytes(self, content: bytes, *, media_type: str) -> str:
  25. assert media_type == "image/png"
  26. self.values.append(content)
  27. return "script-build://raw-artifacts/sha256/" + "a" * 64
  28. @pytest.mark.asyncio
  29. async def test_http_retrieval_revalidates_redirect_and_freezes_response_metadata() -> None:
  30. calls: list[str] = []
  31. def handler(request: httpx.Request) -> httpx.Response:
  32. calls.append(str(request.url))
  33. return httpx.Response(
  34. 200,
  35. json={"data": [{"channel_content_id": "post-1", "title": "case"}]},
  36. )
  37. async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
  38. safe = SafeHttpClient(
  39. client,
  40. OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
  41. )
  42. result = await ExternalRetrievalAdapter("https://api.example.com/search", safe).retrieve(
  43. query={"keyword": "case"}, snapshot=SimpleNamespace()
  44. )
  45. assert calls == ["https://api.example.com/search"]
  46. assert result.source_refs == ("external:post-1",)
  47. assert str(result.metadata["response_sha256"]).startswith("sha256:")
  48. def redirect(_: httpx.Request) -> httpx.Response:
  49. return httpx.Response(302, headers={"location": "https://127.0.0.1/private"})
  50. async with httpx.AsyncClient(transport=httpx.MockTransport(redirect)) as client:
  51. safe = SafeHttpClient(
  52. client,
  53. OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
  54. )
  55. with pytest.raises(UnsafeOutboundTarget):
  56. await safe.request("GET", "https://api.example.com/start")
  57. def oversized(_: httpx.Request) -> httpx.Response:
  58. return httpx.Response(200, content=b"123456")
  59. async with httpx.AsyncClient(transport=httpx.MockTransport(oversized)) as client:
  60. safe = SafeHttpClient(
  61. client,
  62. OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
  63. max_response_bytes=5,
  64. )
  65. with pytest.raises(ProtocolViolation, match="byte limit"):
  66. await safe.request("GET", "https://api.example.com/large")
  67. @pytest.mark.asyncio
  68. async def test_external_response_secrets_are_redacted_before_evidence_summary() -> None:
  69. def handler(_request: httpx.Request) -> httpx.Response:
  70. return httpx.Response(
  71. 200,
  72. json={
  73. "data": [
  74. {
  75. "id": "p1",
  76. "authorization": "Bearer raw-token",
  77. "url": "https://source.example/item?token=secret",
  78. }
  79. ]
  80. },
  81. )
  82. async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
  83. safe = SafeHttpClient(
  84. client,
  85. OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
  86. )
  87. result = await ExternalRetrievalAdapter("https://api.example.com/search", safe).retrieve(
  88. query={"keyword": "case"}, snapshot=SimpleNamespace()
  89. )
  90. assert "raw-token" not in result.summary
  91. assert "secret" not in result.summary
  92. @pytest.mark.asyncio
  93. async def test_pattern_polling_and_image_limits_use_mock_transport() -> None:
  94. poll_count = 0
  95. def handler(request: httpx.Request) -> httpx.Response:
  96. nonlocal poll_count
  97. if request.url.path == "/pattern":
  98. return httpx.Response(
  99. 200,
  100. json={
  101. "status": "pending",
  102. "session_id": "task-1",
  103. "poll_url": "https://api.example.com/poll/task-1",
  104. },
  105. )
  106. if request.url.path.startswith("/poll"):
  107. poll_count += 1
  108. return httpx.Response(200, json={"status": "done", "results": [{"id": 2}]})
  109. return httpx.Response(
  110. 200,
  111. content=b"\x89PNG\r\n\x1a\nimage",
  112. headers={"content-type": "image/png"},
  113. )
  114. async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
  115. safe = SafeHttpClient(
  116. client,
  117. OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
  118. )
  119. result = await PatternRetrievalAdapter(
  120. "https://api.example.com/pattern",
  121. safe,
  122. poll_interval_seconds=0.001,
  123. ).retrieve(query={"message": "shape"}, snapshot=SimpleNamespace())
  124. raw_store = _RawArtifacts()
  125. images = await SafeImageAdapter(
  126. safe,
  127. raw_store,
  128. max_images=1,
  129. max_image_bytes=20,
  130. max_total_image_bytes=20,
  131. ).load(urls=["https://api.example.com/image"], snapshot=SimpleNamespace())
  132. assert poll_count == 1
  133. assert result.source_refs == ("pattern:2",)
  134. assert images[0]["digest"].startswith("sha256:")
  135. assert "url" not in images[0]
  136. assert images[0]["raw_artifact_ref"].startswith("script-build://raw-artifacts/")
  137. assert len(raw_store.values) == 1
  138. @pytest.mark.asyncio
  139. async def test_pattern_timeout_is_frozen_as_a_limitation() -> None:
  140. def handler(_request: httpx.Request) -> httpx.Response:
  141. return httpx.Response(200, json={"status": "pending", "session_id": "task-1"})
  142. async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
  143. safe = SafeHttpClient(
  144. client,
  145. OutboundPolicy(frozenset({"api.example.com"}), resolver=_public_resolver),
  146. )
  147. result = await PatternRetrievalAdapter(
  148. "https://api.example.com/pattern",
  149. safe,
  150. poll_interval_seconds=0.001,
  151. poll_timeout_seconds=0,
  152. ).retrieve(query={"message": "shape"}, snapshot=SimpleNamespace())
  153. assert result.source_refs == ()
  154. assert result.limitations == ("timeout",)
  155. assert result.metadata["task_id"] == "task-1"
  156. @pytest.mark.asyncio
  157. async def test_file_decode_and_knowledge_freeze_file_hash_and_rank(tmp_path) -> None:
  158. decode_path = tmp_path / "decode.json"
  159. decode_path.write_text(
  160. json.dumps({"items": [{"post_id": 4, "account": "acct", "text": "focus"}]}),
  161. encoding="utf-8",
  162. )
  163. knowledge_path = tmp_path / "knowledge.json"
  164. knowledge_path.write_text(json.dumps([{"id": "k1", "title": "focus method"}]), encoding="utf-8")
  165. decode_adapter = FileDecodeRetrievalAdapter(decode_path)
  166. snapshot = SimpleNamespace(
  167. model_manifest={"embedding_model": "frozen-embed"},
  168. datasource_manifest={},
  169. )
  170. decode = await decode_adapter.retrieve(
  171. query={
  172. "keyword": "focus",
  173. "account_name": "acct",
  174. "top_k": 3,
  175. "return_field": "主脉络",
  176. },
  177. snapshot=snapshot,
  178. )
  179. knowledge = await FileKnowledgeRetrievalAdapter(knowledge_path).retrieve(
  180. query={"keyword": "focus", "max_count": 3}, snapshot=snapshot
  181. )
  182. assert decode.metadata["embedding_model"] == "frozen-embed"
  183. assert decode.metadata["ranks"] == [1]
  184. assert decode.metadata["scores"] == [1.0]
  185. assert str(knowledge.metadata["file_sha256"]).startswith("sha256:")
  186. frozen_digest = decode.metadata["index_sha256"]
  187. decode_path.write_text(json.dumps({"items": []}), encoding="utf-8")
  188. replay = await decode_adapter.retrieve(
  189. query={"keyword": "focus", "top_k": 3, "return_field": "主脉络"},
  190. snapshot=SimpleNamespace(
  191. model_manifest={},
  192. datasource_manifest={"decode_index": {"sha256": frozen_digest}},
  193. ),
  194. )
  195. assert replay.source_refs == ("decode:4",)
  196. with pytest.raises(ProtocolViolation, match="changed"):
  197. await FileDecodeRetrievalAdapter(decode_path).retrieve(
  198. query={"keyword": "focus", "top_k": 3, "return_field": "主脉络"},
  199. snapshot=SimpleNamespace(
  200. model_manifest={},
  201. datasource_manifest={"decode_index": {"sha256": frozen_digest}},
  202. ),
  203. )
  204. @pytest.mark.asyncio
  205. async def test_http_decode_preserves_legacy_query_and_ranked_scores() -> None:
  206. requests: list[dict[str, object]] = []
  207. def handler(request: httpx.Request) -> httpx.Response:
  208. requests.append(json.loads(request.content))
  209. return httpx.Response(
  210. 200,
  211. json={"items": [{"post_id": "p1", "score": 0.91, "text": "safe"}]},
  212. )
  213. async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
  214. safe = SafeHttpClient(
  215. client,
  216. OutboundPolicy(frozenset({"decode.example"}), resolver=_public_resolver),
  217. )
  218. result = await HttpDecodeRetrievalAdapter(
  219. "https://decode.example/search",
  220. safe,
  221. ).retrieve(
  222. query={
  223. "return_field": "主脉络",
  224. "keyword": "focus",
  225. "account_name": "acct",
  226. "match_fields": {},
  227. "top_k": 3,
  228. },
  229. snapshot=SimpleNamespace(
  230. datasource_manifest={"decode_index": {"sha256": "sha256:" + "1" * 64}},
  231. model_manifest={"embedding_model": "embed-v1"},
  232. ),
  233. )
  234. assert requests[0]["top_k"] == 3
  235. assert result.source_refs == ("decode:p1",)
  236. assert result.metadata["scores"] == [0.91]
  237. assert result.metadata["ranks"] == [1]
  238. @pytest.mark.asyncio
  239. async def test_uploaded_topic_parser_keeps_old_shape_without_writing() -> None:
  240. gateway = SqlUploadedTopicGateway(SimpleNamespace()) # type: ignore[arg-type]
  241. parsed = await gateway.parse(
  242. {
  243. "选题融合": "topic",
  244. "target_post": {"channel_account_name": "acct"},
  245. "灵感点": [
  246. {
  247. "点": "point",
  248. "实质": {"具体元素": [{"名称": "element", "说明": "note"}]},
  249. }
  250. ],
  251. }
  252. )
  253. assert parsed["account_name"] == "acct"
  254. assert parsed["item_count"] == 1
  255. assert parsed["points"][0]["items"][0]["dimension"] == "实质"
  256. @pytest.mark.asyncio
  257. async def test_uploaded_topic_create_writes_a_readable_legacy_graph(database) -> None:
  258. _, sessions = database
  259. gateway = SqlUploadedTopicGateway(sessions)
  260. parsed = await gateway.parse(
  261. {
  262. "选题融合": "uploaded topic",
  263. "target_post": {"channel_account_name": "acct"},
  264. "关键点": [
  265. {
  266. "点": "point",
  267. "形式": [{"名称": "element", "说明": "note"}],
  268. }
  269. ],
  270. }
  271. )
  272. created = await gateway.create(parsed, account_name=None)
  273. graph = await LegacySqlAlchemyInputReader(sessions).read_topic_graph(
  274. execution_id=0,
  275. topic_build_id=created["topic_build_id"],
  276. topic_id=created["topic_id"],
  277. )
  278. assert graph["execution"] == {"id": 0, "status": "upload"}
  279. assert graph["topic"]["result"] == "uploaded topic"
  280. assert graph["points"][0]["item_ids"] == [graph["composition_items"][0]["id"]]