103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""extractor DEDUP 路径：多条 memory embed 必须并行，并记录 embed 耗时 / meta。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Any, Dict, List
|
|
|
|
from hy_memory.pipelines.base import WriteRequest
|
|
from hy_memory.pipelines.writer import MemoryWriter
|
|
|
|
|
|
def run(coro):
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
return loop.run_until_complete(coro)
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
class _FakeCache:
|
|
def __init__(self):
|
|
self.logs: List[Dict[str, Any]] = []
|
|
|
|
async def store_pipeline_log(self, **kwargs):
|
|
self.logs.append(kwargs)
|
|
|
|
|
|
class TestDedupExtractedParallelEmbed:
|
|
def test_dedup_extracted_uses_write_batch_queue(self):
|
|
"""N 条文本并发进入 write batch queue，并把向量交给后续阶段。"""
|
|
|
|
async def _test():
|
|
delay = 0.08
|
|
in_flight = 0
|
|
max_in_flight = 0
|
|
call_order: List[str] = []
|
|
|
|
class _FakeEmbed:
|
|
async def embed_queued_with_meta(self, text: str):
|
|
nonlocal in_flight, max_in_flight
|
|
call_order.append(text)
|
|
in_flight += 1
|
|
max_in_flight = max(max_in_flight, in_flight)
|
|
await asyncio.sleep(delay)
|
|
in_flight -= 1
|
|
idx = int(text.split("-")[1])
|
|
vec = [0.0] * 8
|
|
vec[idx % 8] = 1.0
|
|
return vec, {
|
|
"http_status": 200,
|
|
"model": "fake-embed",
|
|
"usage": {"prompt_tokens": 3, "total_tokens": 3},
|
|
"response_headers": {"x-request-id": f"req-{idx}"},
|
|
"cache_hit": False,
|
|
}
|
|
|
|
cache = _FakeCache()
|
|
writer = MemoryWriter.__new__(MemoryWriter)
|
|
writer._embed_service = _FakeEmbed() # type: ignore[assignment]
|
|
writer._cache = cache
|
|
|
|
texts = [f"memory-{i}-content-about-topic-{i}" for i in range(6)]
|
|
metas: List[Dict[str, Any]] = [
|
|
{"content": t, "layer": "L2_FACT", "tags": []} for t in texts
|
|
]
|
|
req = WriteRequest(
|
|
content="irrelevant",
|
|
user_id="u1",
|
|
agent_id="a1",
|
|
)
|
|
|
|
t0 = asyncio.get_event_loop().time()
|
|
out_texts, out_metas, out_embeddings = await writer._dedup_extracted(
|
|
texts, metas, req, req_id="rq-parallel",
|
|
)
|
|
wall = asyncio.get_event_loop().time() - t0
|
|
|
|
assert out_texts == texts
|
|
assert len(out_metas) == 6
|
|
assert set(out_embeddings) == set(texts)
|
|
assert set(call_order) == set(texts)
|
|
assert max_in_flight >= 2
|
|
assert wall < 0.30, f"expected parallel wall <0.30s, got {wall:.3f}s"
|
|
|
|
assert len(cache.logs) == 1
|
|
parsed = json.loads(cache.logs[0]["parsed"])
|
|
assert parsed["trigger"] == "extractor"
|
|
assert parsed["embedding_source"] == "embed_service"
|
|
assert parsed["embed_count"] == 6
|
|
assert parsed["embed_ok_count"] == 6
|
|
assert isinstance(parsed["embed_elapsed_ms"], (int, float))
|
|
assert parsed["embed_elapsed_ms"] < 300
|
|
assert len(parsed["embed_metas"]) == 6
|
|
for m in parsed["embed_metas"]:
|
|
assert "embedding" not in m
|
|
assert m.get("model") == "fake-embed"
|
|
assert "elapsed_ms" in m
|
|
assert "response_headers" in m
|
|
|
|
run(_test())
|