97 lines
2.6 KiB
Python
97 lines
2.6 KiB
Python
|
from uuid import uuid4
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from app.core.errors import RecoverableAgentError
|
||
|
|
from app.infrastructure.milvus_profile_projection import MilvusProfileProjection
|
||
|
|
|
||
|
|
|
||
|
|
class FakeMilvus:
|
||
|
|
def __init__(self, existing: list[dict[str, object]] | None = None) -> None:
|
||
|
|
self.existing = existing or []
|
||
|
|
self.queries: list[dict[str, object]] = []
|
||
|
|
self.upserts: list[dict[str, object]] = []
|
||
|
|
|
||
|
|
async def query(self, **kwargs: object) -> list[dict[str, object]]:
|
||
|
|
self.queries.append(kwargs)
|
||
|
|
return self.existing
|
||
|
|
|
||
|
|
async def upsert(self, **kwargs: object) -> None:
|
||
|
|
self.upserts.append(kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
def payload() -> dict[str, object]:
|
||
|
|
return {
|
||
|
|
"customer_id": 7,
|
||
|
|
"profile_version": 1,
|
||
|
|
"memory_sources": [{
|
||
|
|
"memory_uuid": str(uuid4()),
|
||
|
|
"memory_key": "preference:risk_level",
|
||
|
|
"content": "稳健型",
|
||
|
|
"memory_type": "preference",
|
||
|
|
"confidence": 0.9,
|
||
|
|
"version": 2,
|
||
|
|
"valid_until": None,
|
||
|
|
}],
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_upsert_writes_schema_fields_and_vector() -> None:
|
||
|
|
client = FakeMilvus()
|
||
|
|
projection = MilvusProfileProjection(client, _embed)
|
||
|
|
|
||
|
|
await projection.upsert(payload())
|
||
|
|
|
||
|
|
assert len(client.upserts) == 1
|
||
|
|
row = client.upserts[0]["data"][0]
|
||
|
|
assert row["customer_id"] == 7
|
||
|
|
assert row["status"] == "active"
|
||
|
|
assert len(row["embedding"]) == 1024
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_lower_memory_version_is_not_overwritten() -> None:
|
||
|
|
data = payload()
|
||
|
|
source = data["memory_sources"][0]
|
||
|
|
assert isinstance(source, dict)
|
||
|
|
memory_uuid = source["memory_uuid"]
|
||
|
|
client = FakeMilvus(existing=[{
|
||
|
|
"memory_uuid": memory_uuid, "customer_id": 7, "version": 3,
|
||
|
|
}])
|
||
|
|
|
||
|
|
await MilvusProfileProjection(client, _embed).upsert(data)
|
||
|
|
|
||
|
|
assert client.upserts == []
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_embedding_dimension_is_enforced() -> None:
|
||
|
|
with pytest.raises(RecoverableAgentError, match="维度"):
|
||
|
|
await MilvusProfileProjection(client=FakeMilvus(), embed=_embed_short).upsert(
|
||
|
|
payload()
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_non_uuid_memory_id_is_rejected() -> None:
|
||
|
|
data = payload()
|
||
|
|
source = data["memory_sources"][0]
|
||
|
|
assert isinstance(source, dict)
|
||
|
|
source["memory_uuid"] = "unsafe\" or true"
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match="memory_uuid"):
|
||
|
|
await MilvusProfileProjection(FakeMilvus(), _embed).upsert(data)
|
||
|
|
|
||
|
|
|
||
|
|
def _vector(size: int = 1024) -> list[float]:
|
||
|
|
return [0.0] * size
|
||
|
|
|
||
|
|
|
||
|
|
async def _embed(_: str) -> list[float]:
|
||
|
|
return _vector()
|
||
|
|
|
||
|
|
|
||
|
|
async def _embed_short(_: str) -> list[float]:
|
||
|
|
return _vector(3)
|